Spaces:
Running on Zero
Running on Zero
File size: 2,261 Bytes
8a28a8d | 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 | import gc
from typing import List
import gradio as gr
from imagegen_utils.app_utils import _ensure_model_downloaded
from core.settings import ALL_MODEL_MAP
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.initialized = True
print("✅ ModelManager initialized.")
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 / max(len(required_models), 1), desc=f"Checking file: {display_name}")
try:
_ensure_model_downloaded(display_name, progress)
except Exception as e:
raise gr.Error(f"模型“{display_name}”下载失败：{e}")
print(f"--- [ModelManager] ✅ All required models are present on disk. ---")
model_manager = ModelManager()
def release_loaded_models() -> bool:
"""Best-effort release of ComfyUI model state while GPU access is active."""
released = False
try:
from comfy import model_management
unload = getattr(model_management, "unload_all_models", None)
if callable(unload):
unload()
released = True
cleanup = getattr(model_management, "cleanup_models", None)
if callable(cleanup):
cleanup()
gc.collect()
empty_cache = getattr(model_management, "soft_empty_cache", None)
if callable(empty_cache):
try:
empty_cache(force=True)
except TypeError:
empty_cache()
print("✅ Released ComfyUI model state after a model switch/error.")
except Exception as exc:
# Cleanup must never hide the original generation result or exception.
gc.collect()
print(f"Warning: Could not fully release ComfyUI model state: {exc}")
return released
|