shunt-adapter-testing / model_manager.py
AbstractPhil
local project created to properly edit and debug
d3479d5
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