Spaces:
Sleeping
Sleeping
File size: 7,094 Bytes
5b29993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
import gc
from typing import Dict, List, Any, Set
import torch
import gradio as gr
from comfy import model_management
from core.settings import ALL_MODEL_MAP, CHECKPOINT_DIR, LORA_DIR, DIFFUSION_MODELS_DIR, VAE_DIR, TEXT_ENCODERS_DIR
from comfy_integration.nodes import checkpointloadersimple, LoraLoader
from nodes import NODE_CLASS_MAPPINGS
from utils.app_utils import get_value_at_index, _ensure_model_downloaded
class ModelManager:
_instance = None
def __new__(cls, *args, **kwargs):
if not cls._instance:
cls._instance = super(ModelManager, cls).__new__(cls, *args, **kwargs)
return cls._instance
def __init__(self):
if hasattr(self, 'initialized'):
return
self.loaded_models: Dict[str, Any] = {}
self.initialized = True
print("✅ ModelManager initialized.")
def get_loaded_model_names(self) -> Set[str]:
return set(self.loaded_models.keys())
def _load_single_model(self, display_name: str, progress) -> Any:
print(f"--- [ModelManager] Loading model: '{display_name}' ---")
filename = _ensure_model_downloaded(display_name, progress)
_, _, model_type, _ = ALL_MODEL_MAP[display_name]
loader_map = {
"SDXL": (checkpointloadersimple, "load_checkpoint", {"ckpt_name": filename}),
"SD1.5": (checkpointloadersimple, "load_checkpoint", {"ckpt_name": filename}),
"UNET": (NODE_CLASS_MAPPINGS["UNETLoader"](), "load_unet", {"unet_name": filename, "weight_dtype": "default"}),
"VAE": (NODE_CLASS_MAPPINGS["VAELoader"](), "load_vae", {"vae_name": filename}),
"TEXT_ENCODER": (NODE_CLASS_MAPPINGS["CLIPLoader"](), "load_clip", {"clip_name": filename, "type": "wan", "device": "default"}),
}
if model_type not in loader_map:
if model_type == "LORA":
print(f"--- [ModelManager] ✅ '{display_name}' is a LoRA. It will be loaded dynamically. ---")
return (filename,)
raise ValueError(f"[ModelManager] No loader configured for model type '{model_type}'")
loader_instance, method_name, kwargs = loader_map[model_type]
load_method = getattr(loader_instance, method_name)
loaded_tuple = load_method(**kwargs)
print(f"--- [ModelManager] ✅ Successfully loaded '{display_name}' to CPU/RAM ---")
return loaded_tuple
def move_models_to_gpu(self, required_models: List[str]):
print(f"--- [ModelManager] Moving models to GPU: {required_models} ---")
models_to_load_gpu = []
for name in required_models:
if name in self.loaded_models:
model_tuple = self.loaded_models[name]
_, _, model_type, _ = ALL_MODEL_MAP[name]
if model_type in ["SDXL", "SD1.5"]:
models_to_load_gpu.append(get_value_at_index(model_tuple, 0))
if models_to_load_gpu:
model_management.load_models_gpu(models_to_load_gpu)
print("--- [ModelManager] ✅ Models successfully moved to GPU. ---")
else:
print("--- [ModelManager] ⚠️ No checkpoint models found to move to GPU. ---")
def ensure_models_downloaded(self, required_models: List[str], progress):
print(f"--- [ModelManager] Ensuring models are downloaded: {required_models} ---")
for i, display_name in enumerate(required_models):
if progress and hasattr(progress, '__call__'):
progress(i / len(required_models), desc=f"Checking file: {display_name}")
try:
_ensure_model_downloaded(display_name, progress)
except Exception as e:
raise gr.Error(f"Failed to download model '{display_name}'. Reason: {e}")
print(f"--- [ModelManager] ✅ All required models are present on disk. ---")
def load_managed_models(self, required_models: List[str], active_loras: List[Dict[str, Any]], progress) -> Dict[str, Any]:
required_set = set(required_models)
current_set = set(self.loaded_models.keys())
loras_changed = len(active_loras) > 0 or len(current_set - required_set) > 0
models_to_unload = current_set - required_set
if models_to_unload or loras_changed:
if models_to_unload:
print(f"--- [ModelManager] Models to unload: {models_to_unload} ---")
if loras_changed and not models_to_unload:
models_to_unload = current_set.intersection(required_set)
print(f"--- [ModelManager] LoRA configuration changed. Reloading base model(s): {models_to_unload} ---")
model_management.unload_all_models()
self.loaded_models.clear()
gc.collect()
torch.cuda.empty_cache()
print("--- [ModelManager] All models unloaded to free RAM. ---")
models_to_load = required_set if (models_to_unload or loras_changed) else required_set - current_set
if models_to_load:
print(f"--- [ModelManager] Models to load: {models_to_load} ---")
for i, display_name in enumerate(models_to_load):
progress(i / len(models_to_load), desc=f"Loading model: {display_name}")
try:
loaded_model_data = self._load_single_model(display_name, progress)
if active_loras and ALL_MODEL_MAP[display_name][2] in ["SDXL", "SD1.5"]:
print(f"--- [ModelManager] Applying {len(active_loras)} LoRAs on CPU... ---")
lora_loader = LoraLoader()
patched_model, patched_clip = loaded_model_data[0], loaded_model_data[1]
for lora_info in active_loras:
patched_model, patched_clip = lora_loader.load_lora(
model=patched_model,
clip=patched_clip,
lora_name=lora_info["lora_name"],
strength_model=lora_info["strength_model"],
strength_clip=lora_info["strength_clip"]
)
loaded_model_data = (patched_model, patched_clip, loaded_model_data[2])
print(f"--- [ModelManager] ✅ All LoRAs merged into the model on CPU. ---")
self.loaded_models[display_name] = loaded_model_data
except Exception as e:
raise gr.Error(f"Failed to load model or apply LoRA '{display_name}'. Reason: {e}")
else:
print(f"--- [ModelManager] All required models are already loaded. ---")
return {name: self.loaded_models[name] for name in required_models}
model_manager = ModelManager() |