from typing import Dict, Optional, Any, Union, Tuple import os import torch import torch.nn as nn import logging from pathlib import Path from dataclasses import dataclass from enum import Enum from safetensors.torch import load_file from torch.nn import Module from transformers import AutoModel, AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM, BertModel, BertTokenizer, \ PreTrainedTokenizerFast, T5TokenizerFast, T5EncoderModel from .custom.t5_encoder_with_projection import T5EncoderWithProjection logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # # Helper for namespaced cache keys def _make_key(model_type: str, model_id: str) -> str: """ Produce a unique key for the internal cache. Example ------- >>> _make_key("bert", "bert-base") 'bert:bert-base' """ return f"{model_type}:{model_id}" # Thread-safe registry wrapper class _SafeDict(dict): """A dict protected by a re-entrant lock for thread-safe writes.""" def __init__(self): super().__init__() import threading self._lock = threading.RLock() def safe_set(self, key, value): with self._lock: super().__setitem__(key, value) def safe_get(self, key, default=None): with self._lock: return super().get(key, default) def safe_del(self, key): with self._lock: if key in self: super().__delitem__(key) return True return False # -------------------------------------------------------------------------------------------------------------------- # # WARNING: ENABLING THIS TRUST_REMOTE_CODE FLAG WILL ALLOW EXECUTION OF ARBITRARY CODE FROM THE MODEL REPOSITORY. # USE WITH EXTREME CAUTION, AS IT CAN POTENTIALLY EXECUTE MALICIOUS CODE FROM UNTRUSTED SOURCES. TRUST_REMOTE_CODE = False # Set to True only if you trust the source of the models you are loading. # I advise leaving this OFF for any production or sensitive environments, and for any government or enterprise use. # Ensure you fully trust the model repository and its maintainers and reviewing the code thoroughly. # You cannot ONLY trust an AI's response to the question of whether it is safe to enable this flag, # as it may not have the full context of security implications or the specific model's behavior. # -------------------------------------------------------------------------------------------------------------------- # # COMFYUI operates within a form of sandbox, but enabling remote code execution can still pose many unseen risks. # -------------------------------------------------------------------------------------------------------------------- # class ModelType(Enum): """Enum for different model types""" SHUNT_ADAPTER = "shunt_adapter" T5_MODEL = "t5_model" BERT_MODEL = "bert" NOMIC_BERT_MODEL = "nomic_bert" GENERIC = "generic" TOKENIZER = "tokenizer" @dataclass class ModelInfo: """Container for model information""" model: nn.Module model_type: ModelType config: Dict[str, Any] device: torch.device dtype: torch.dtype metadata: Dict[str, Any] = None trust_remote_code: bool = TRUST_REMOTE_CODE # Use global setting by default class ModelManager: """ Centralized model loader / cache with thread-safety and namespaced keys. """ def __init__(self, cache_dir: Optional[str] = None): # Thread-safe model cache self.models: _SafeDict = _SafeDict() self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.cache_dir = self._setup_cache_dir(cache_dir) # be VERY careful with huggingface keys, remote code execution, and model downloads. # If you are using private models or need to authenticate, set the HuggingFace API key. def set_huggingface_key(self, key: str): """ Set the HuggingFace API key for model downloads. This is useful if you have a private model or need to authenticate. """ os.environ["HF_TOKEN"] = key logger.info("HuggingFace API key set successfully.") def get_huggingface_key(self) -> Optional[str]: """ Get the HuggingFace API key if set. This is useful for debugging or checking if authentication is needed. """ return os.environ.get("HF_TOKEN") def set_huggingface_cache_directory(self, directory: str): """ Set the cache directory for HuggingFace model downloads. This is useful if you want to change the cache location. This will not move your models, it only sets the new default directory. """ os.environ["HF_HOME"] = directory logger.info(f"HuggingFace default directory set to: {directory}") def get_huggingface_cache_directory(self) -> Optional[str]: """ Get the cache directory for HuggingFace model downloads. This is useful for debugging or checking where models are stored. """ return os.environ.get("HF_HOME", str(self.cache_dir)) # --------------------------------------------------------------------- # # Internal helpers def _store(self, key: str, info: "ModelInfo") -> None: """Thread-safe insertion into the model cache.""" self.models.safe_set(key, info) def _setup_cache_dir(self, cache_dir: Optional[str]) -> Path: """Setup and validate cache directory""" if cache_dir: cache_path = Path(cache_dir) else: # Use default HuggingFace cache location cache_path = Path.home() / ".cache" / "huggingface" / "transformers" cache_path.mkdir(parents=True, exist_ok=True) logger.info(f"Using cache directory: {cache_path}") return cache_path def get_model(self, key: str) -> Optional["ModelInfo"]: """Retrieve a model by its namespaced key.""" return self.models.safe_get(key) def is_loaded(self, key: str) -> bool: """Return True if the namespaced key is present in the cache.""" return self.models.safe_get(key) is not None def move_model( self, namespaced_key: str, *, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ) -> Optional[nn.Module]: """ Convert device/dtype of a cached model and return the updated object. """ model = self._maybe_convert_dtype(namespaced_key, dtype, device) if model is None: logger.warning("move_model: %s not found", namespaced_key) return model def load_tokenizer( self, id: str, tokenizer_name_or_path: str, target_output_device: Optional[torch.device] = None, force_reload: bool = False, trust_remote_code: Optional[bool] = None, ) -> Optional[tuple[PreTrainedTokenizerFast, dict[str, Any]]]: """Load or fetch from cache a Hugging-Face tokenizer.""" key = _make_key("tokenizer", id) if not force_reload and self.is_loaded(key): model_info = self.get_model(key) return model_info.model, model_info.metadata try: trust_remote_code = ( trust_remote_code if trust_remote_code is not None else TRUST_REMOTE_CODE ) tok = AutoTokenizer.from_pretrained( tokenizer_name_or_path, trust_remote_code=trust_remote_code ) self._store( key, ModelInfo( model=tok, model_type=ModelType.TOKENIZER, config={"tokenizer_name": tokenizer_name_or_path}, device=target_output_device or torch.device("cpu"), dtype=torch.float32, metadata={"source": "huggingface", "trust_remote_code": trust_remote_code}, ), ) logger.info("Loaded tokenizer %s", key) return tok, self.get_model(key).metadata except Exception: logger.exception("Failed to load tokenizer %s", id) return None def load_shunt_adapter( self, adapter_id: str, config: Dict[str, Any], path: Optional[str] = None, repo_id: Optional[str] = None, filename: Optional[str] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, force_reload: bool = False ) -> Optional[nn.Module]: """ Load a shunt adapter from local path or HuggingFace. Args: adapter_id: Unique identifier for the adapter config: Configuration dictionary for the adapter path: Local path to the adapter file repo_id: HuggingFace repository ID filename: Filename in the HuggingFace repository device: Target device dtype: Target dtype force_reload: Force reload even if cached Returns: Loaded adapter model or None if failed """ if not force_reload and self.is_loaded(adapter_id): logger.info(f"Using cached adapter: {adapter_id}") return self._maybe_convert_dtype(adapter_id, dtype, device) try: # Import here to avoid circular imports from two_stream_shunt_adapter import ConditionModulationShuntAdapter # Determine file location file_path = self._resolve_file_path(path, repo_id, filename) if not file_path: raise FileNotFoundError(f"Could not find adapter file for {adapter_id}") # Initialize adapter # if the filename ends with t5-vit-l-14-dual_shunt_booru_13_000_000.safetensors we set attention heads to 4, else we set to 12 logger.info(f"Loading adapter {adapter_id} from {file_path}") adapter = ConditionModulationShuntAdapter(config=config) logger.info(f"Initialized adapter {adapter_id} with config: {config}") # Load weights state_dict = load_file(file_path) logger.info(f"Loaded state_dict for adapter {adapter_id} from {file_path}") adapter.load_state_dict(state_dict, strict=False) logger.info(f"Adapter {adapter_id} state_dict loaded successfully") # Move to device and dtype device = device or self.device dtype = dtype or torch.float32 logger.info(f"Moving adapter {adapter_id} to device: {device}, dtype: {dtype}") adapter = adapter.to(device=device, dtype=dtype) logger.info(f"Adapter {adapter_id} moved to device and dtype successfully") # Cache the model self.models[adapter_id] = ModelInfo( model=adapter, model_type=ModelType.SHUNT_ADAPTER, config=config, device=device, dtype=dtype, metadata={"file_path": str(file_path)} ) logger.info(f"Adapter {adapter_id} cached successfully") logger.info(f"Successfully loaded adapter: {adapter_id}") return adapter except Exception as e: logger.error(f"Failed to load adapter {adapter_id} from {path or repo_id}/{filename}: {e}") logger.debug(f"Traceback: {e.__traceback__}") return None def load_encoder_model(self, model_type: str, # use this to see if it's compatible with the current model manager model_id: str, model_name_or_path: str, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, force_reload: bool = False, trust_remote_code: Optional[bool] = None, # Overrides the global TRUST_REMOTE_CODE setting. config: Optional[Dict[str, Any]] = None # Additional configuration for the model ) -> Optional[nn.Module]: """ Load an encoder model (e.g., BERT, T5) and return it. Args: model_type: Type of the model (e.g., "bert", "t5") model_id: Unique identifier for the model model_name_or_path: Model name or path device: Target device dtype: Target dtype force_reload: Force reload even if cached Returns: Loaded model or None if failed """ if model_type == "bert": return self.load_bert_model(model_id, model_name_or_path, device, dtype, force_reload, trust_remote_code) elif model_type == "nomic_bert": # Nomic BERT is a specific variant of BERT, so we can use the same loading function return self.load_bert_model(model_id, model_name_or_path, device, dtype, force_reload, trust_remote_code) elif "t5" in model_type: return self.load_t5_model(model_id, model_name_or_path, device, dtype, force_reload, trust_remote_code, config) else: logger.error(f"Unsupported model type: {model_type}") return None def load_bert_model( self, model_id: str, model_name_or_path: str, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, force_reload: bool = False, trust_remote_code: Optional[bool] = None # Overrides the global TRUST_REMOTE_CODE setting. ) -> Optional[Tuple[nn.Module, Any]]: """ Load a BERT model and tokenizer. Returns: Tuple of (model, tokenizer) or None if failed """ if not force_reload and self.is_loaded(model_id): logger.info(f"Using cached BERT model: {model_id}") model_info = self.get_model(model_id) return model_info.model, model_info.metadata.get("tokenizer") try: device = device or self.device dtype = dtype or torch.float32 config = AutoConfig.from_pretrained( model_name_or_path, trust_remote_code=trust_remote_code if trust_remote_code is not None else TRUST_REMOTE_CODE # Use the global flag for remote code execution ) # Load tokenizer and model tokenizer = AutoTokenizer.from_pretrained( model_name_or_path, config=config, use_special_tokens=True, # Ensure special tokens are used trust_remote_code=trust_remote_code if trust_remote_code is not None else TRUST_REMOTE_CODE # Use the global flag for remote code execution ) model = AutoModel.from_pretrained( model_name_or_path, config=config, torch_dtype=dtype, trust_remote_code=trust_remote_code if trust_remote_code is not None else TRUST_REMOTE_CODE # Use the global flag for remote code execution ).to(device) # Cache the model self._store(_make_key("bert", model_id), ModelInfo( model=model, model_type=ModelType.BERT_MODEL, config={"model_name": model_name_or_path}, device=device, dtype=dtype, metadata={"tokenizer": tokenizer}, trust_remote_code=trust_remote_code if trust_remote_code is not None else TRUST_REMOTE_CODE )) logger.info(f"Successfully loaded BERT model: {model_id}") return model, tokenizer except Exception as e: logger.error(f"Failed to load BERT model {model_id}: {e}") return None def load_t5_model( self, model_id: str, model_name_or_path: str, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, force_reload: bool = False, override_remote_code: Optional[bool] = None, # Overrides the global TRUST_REMOTE_CODE setting. config: Optional[Dict[str, Any]] = None # Additional configuration for the model ) -> Optional[Tuple[nn.Module, Any]]: """ Load a T5 model and tokenizer. Returns: Tuple of (model, tokenizer) or None if failed """ if not force_reload and self.is_loaded(model_id): logger.info(f"Using cached T5 model: {model_id}") model_info = self.get_model(model_id) return model_info.model, model_info.metadata.get("tokenizer") try: device = device or self.device dtype = dtype or torch.float32 trust_remote_code = override_remote_code if override_remote_code is not None else TRUST_REMOTE_CODE # Load tokenizer and model if config.get("type", "t5") == "t5": tokenizer = AutoTokenizer.from_pretrained( "google/flan-t5-base", trust_remote_code=trust_remote_code # Use the global flag for remote code execution ) elif config.get("type", "t5") == "t5_unchained": tokenizer = T5TokenizerFast.from_pretrained( "AbstractPhil/t5xxl-unchained", trust_remote_code=trust_remote_code # Use the global flag for remote code execution ) else: tokenizer = T5TokenizerFast.from_pretrained( "google/flan-t5-base", trust_remote_code=trust_remote_code # Use the global flag for remote code execution ) if config.get("type", "t5") == "t5": logger.info(f"Loading T5ForConditionalGeneration model from {model_name_or_path}") model = AutoModelForSeq2SeqLM.from_pretrained( model_name_or_path, torch_dtype=dtype, trust_remote_code=trust_remote_code # Use the global flag for remote code execution ).to(device) elif config.get("type", "t5") == "t5_encoder_with_projection": # Load T5EncoderModel with projection layer logger.info(f"Loading T5EncoderWithProjection model from {model_name_or_path}") model = T5EncoderWithProjection.from_pretrained( model_name_or_path, torch_dtype=dtype, trust_remote_code=trust_remote_code # Use the global flag for remote code execution ).to(device) else: # Load standard T5 model logger.info(f"Loading T5EncoderModel from {model_name_or_path}") model = AutoModel.from_pretrained( model_name_or_path, torch_dtype=dtype, trust_remote_code=trust_remote_code # Use the global flag for remote code execution ).to(device) # Cache the model self._store(_make_key("t5", model_id), ModelInfo( model=model, model_type=ModelType.T5_MODEL, config={"model_name": model_name_or_path}, device=device, dtype=dtype, metadata={"tokenizer": tokenizer} )) logger.info(f"Successfully loaded T5 model: {model_id}") return model, tokenizer except Exception as e: logger.error(f"Failed to load T5 model {model_id}: {e}") return None def unload_model(self, model_id: str) -> bool: """ Unload a model to free memory. Returns: True if successfully unloaded, False otherwise """ if model_id in self.models: try: # Move to CPU first to free GPU memory model_info = self.models[model_id] model_info.model.cpu() # Delete the model del self.models[model_id] # Force garbage collection import gc gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info(f"Successfully unloaded model: {model_id}") return True except Exception as e: logger.error(f"Failed to unload model {model_id}: {e}") return False else: logger.warning(f"Model {model_id} not found in cache") return False def list_models(self) -> Dict[str, Dict[str, Any]]: """List all loaded models with their information""" return { model_id: { "type": info.model_type.value, "device": str(info.device), "dtype": str(info.dtype), "config": info.config } for model_id, info in self.models.items() } def clear_all(self): """Clear all loaded models""" model_ids = list(self.models.keys()) for model_id in model_ids: self.unload_model(model_id) logger.info("All models cleared from memory") def _resolve_file_path( self, local_path: Optional[str], repo_id: Optional[str], filename: Optional[str] ) -> Optional[Path]: """Resolve file path from local or HuggingFace""" # Try local path first if local_path and os.path.exists(local_path): return Path(local_path) # Try HuggingFace if repo_id and filename: try: from huggingface_hub import hf_hub_download file_path = hf_hub_download( repo_id=repo_id, filename=filename, cache_dir=str(self.cache_dir), repo_type="model" ) return Path(file_path) except Exception as e: logger.error(f"Failed to download from HuggingFace: {e}") return None def _maybe_convert_dtype( self, model_id: str, target_dtype: Optional[torch.dtype], target_device: Optional[torch.device] ) -> Optional[nn.Module]: """Convert model dtype/device if needed""" model_info = self.get_model(model_id) if not model_info: return None model = model_info.model changed = False # Check dtype conversion if target_dtype and model_info.dtype != target_dtype: try: model = model.to(dtype=target_dtype) model_info.dtype = target_dtype changed = True logger.info(f"Converted {model_id} to dtype: {target_dtype}") except Exception as e: logger.error(f"Failed to convert dtype for {model_id}: {e}") # Check device conversion if target_device and model_info.device != target_device: try: model = model.to(device=target_device) model_info.device = target_device changed = True logger.info(f"Moved {model_id} to device: {target_device}") except Exception as e: logger.error(f"Failed to move {model_id} to device: {e}") if changed: model_info.model = model return model def __del__(self): """Cleanup on deletion""" self.clear_all() # Global instance (singleton pattern) _global_model_manager: Optional[ModelManager] = None def get_model_manager(cache_dir: Optional[str] = None) -> ModelManager: """Get or create the global model manager instance""" global _global_model_manager if _global_model_manager is None: _global_model_manager = ModelManager(cache_dir=cache_dir) return _global_model_manager