""" ComfyUI TTS Model Manager """ import weakref import gc from typing import Optional, Any, Dict from .base_wrapper import ComfyUIModelWrapper, ModelInfo, COMFYUI_AVAILABLE, model_management class ComfyUITTSModelManager: """ Manager that integrates TTS models with ComfyUI's model management system. This replaces static caches with ComfyUI-managed model loading/unloading. """ def __init__(self): self._model_cache: Dict[str, ComfyUIModelWrapper] = {} def _offload_conflicting_tts_models(self, engine: str, target_device: str, keep_model_key: Optional[str] = None) -> int: """ Move conflicting GPU-resident TTS models to CPU before loading/moving another one. This keeps fresh load_model() behavior aligned with ensure_device(): if a new TTS model wants CUDA, other cached TTS engines should not stay in VRAM. """ if not str(target_device).startswith('cuda'): return 0 models_to_offload = [] for cache_key, wrapper in list(self._model_cache.items()): should_offload = ( wrapper.model_info.model_type == "tts" and wrapper._is_loaded_on_gpu and (wrapper.model_info.engine != engine or (keep_model_key is not None and cache_key != keep_model_key)) ) if should_offload: models_to_offload.append((cache_key, wrapper)) if not models_to_offload: return 0 print(f"๐Ÿ—‘๏ธ Moving {len(models_to_offload)} TTS model(s) to CPU to make room for {engine}") for cache_key, wrapper in models_to_offload: try: wrapper.partially_unload('cpu', wrapper._memory_size) except Exception as e: print(f"โš ๏ธ Failed to move {wrapper.model_info.engine} to CPU, removing instead: {e}") self.remove_model(cache_key) import torch if torch.cuda.is_available(): torch.cuda.synchronize() torch.cuda.empty_cache() import time time.sleep(0.1) return len(models_to_offload) def load_model(self, model_factory_func, model_key: str, model_type: str, engine: str, device: str, force_reload: bool = False, **factory_kwargs) -> ComfyUIModelWrapper: """ Load a model using ComfyUI's model management system. Args: model_factory_func: Function that creates the model model_key: Unique key for caching model_type: Type of model ("tts", "vc", etc.) engine: Engine name ("chatterbox", "f5tts", etc.) device: Target device **factory_kwargs: Arguments for model factory function Returns: ComfyUI-wrapped model """ # No more shadow storage - if model was destroyed, create completely fresh # Check if already cached if model_key in self._model_cache and not force_reload: wrapper = self._model_cache[model_key] is_valid = getattr(wrapper, '_is_valid_for_reuse', True) print(f"๐Ÿ” Cache check for {model_type} ({engine}): valid={is_valid}, force_reload={force_reload}") # Check if cached model is still valid for reuse if not is_valid: # For Higgs Audio with CUDA graph corruption, try to reinitialize in-place if engine == "higgs_audio": print(f"๐Ÿ”„ Attempting in-place reinitializion of corrupted {engine} model to avoid memory conflicts") try: # Reset CUDA graph state without creating new model if hasattr(wrapper.model, 'engine') and hasattr(wrapper.model.engine, 'cuda_graphs_initialized'): wrapper.model.engine.cuda_graphs_initialized = False print(f"โœ… Reset CUDA graph state for existing model") # Move back to GPU for reinit wrapper.model_load(device) # Mark as valid again wrapper._is_valid_for_reuse = True print(f"โœ… Successfully reinitialized {engine} model in-place") return wrapper except Exception as e: print(f"โš ๏ธ In-place reinit failed: {e}, falling back to full recreation") # For VibeVoice, try to reinitialize corrupted model state # Unlike Higgs Audio, VibeVoice doesn't use CUDA graphs so should be recoverable elif engine == "vibevoice": print(f"๐Ÿ”„ VibeVoice: Attempting to recover from CPU offloading corruption") try: # Clear any cached internal state that might be corrupted if hasattr(wrapper.model, '_past_key_values'): wrapper.model._past_key_values = None if hasattr(wrapper.model, '_cache'): wrapper.model._cache = None # Reset model to evaluation mode and clear gradients wrapper.model.eval() if hasattr(wrapper.model, 'zero_grad'): wrapper.model.zero_grad() # Move back to GPU with proper state reset wrapper.model_load(device) # Mark as valid again wrapper._is_valid_for_reuse = True print(f"โœ… Successfully recovered VibeVoice model from corruption") return wrapper except Exception as e: print(f"โš ๏ธ VibeVoice recovery failed: {e}, falling back to full recreation") # For IndexTTS-2, try to recover from device mismatch after CPU offloading elif engine == "index_tts": print(f"๐Ÿ”„ IndexTTS-2: Attempting to recover from CPU offloading device mismatch") try: # IndexTTS-2 has multiple model components that need device synchronization # Clear any device-cached state if hasattr(wrapper.model, '_model_config'): wrapper.model._model_config = None # Reset model to evaluation mode if hasattr(wrapper.model, 'eval'): wrapper.model.eval() # Force device reload for all model components wrapper.model_load(device) # Mark as valid again wrapper._is_valid_for_reuse = True print(f"โœ… Successfully recovered IndexTTS-2 model from device mismatch") return wrapper except Exception as e: print(f"โš ๏ธ IndexTTS-2 recovery failed: {e}, falling back to full recreation") print(f"๐Ÿ—‘๏ธ Removing invalid cached model: {model_type} ({engine}) - corrupted by previous unload") self.remove_model(model_key) # Continue to create new model below else: print(f"โ™ป๏ธ Reusing valid cached model: {model_type} ({engine})") # Before loading to GPU, clear other models to make room if wrapper.current_device != device and device != 'auto' and device.startswith('cuda'): # Clear models from different engines to free VRAM cached_models = list(self._model_cache.keys()) models_to_clear = [] for cache_key in cached_models: cached_wrapper = self._model_cache[cache_key] # Clear TTS models from different engines that are on GPU if (cached_wrapper.model_info.engine != engine and cached_wrapper.model_info.model_type == "tts" and cached_wrapper._is_loaded_on_gpu): models_to_clear.append(cache_key) if models_to_clear: print(f"๐Ÿ—‘๏ธ Clearing {len(models_to_clear)} TTS models to free VRAM for reused model") for key in models_to_clear: self.remove_model(key) # Ensure model is loaded on correct device if wrapper.current_device != device and device != 'auto': wrapper.model_load(device) return wrapper elif force_reload and model_key in self._model_cache: wrapper = self._model_cache[model_key] # For Higgs Audio, try in-place reinitialization instead of full recreation if engine == "higgs_audio": print(f"๐Ÿ”„ Force reload: attempting in-place reinitializion of {engine} model to avoid memory conflicts") try: # Reset CUDA graph state without creating new model if hasattr(wrapper.model, 'engine') and hasattr(wrapper.model.engine, 'cuda_graphs_initialized'): wrapper.model.engine.cuda_graphs_initialized = False print(f"โœ… Reset CUDA graph state for existing model") # Move back to GPU for reinit wrapper.model_load(device) # Mark as valid again wrapper._is_valid_for_reuse = True print(f"โœ… Successfully reinitialized {engine} model in-place (force reload)") return wrapper except Exception as e: print(f"โš ๏ธ Force reload in-place reinit failed: {e}, falling back to full recreation") # For IndexTTS-2, try in-place device synchronization on force reload elif engine == "index_tts": print(f"๐Ÿ”„ Force reload: attempting IndexTTS-2 device synchronization") try: # Clear device-cached state if hasattr(wrapper.model, '_model_config'): wrapper.model._model_config = None # Force device reload for all model components wrapper.model_load(device) # Mark as valid again wrapper._is_valid_for_reuse = True print(f"โœ… Successfully reloaded IndexTTS-2 model with device sync (force reload)") return wrapper except Exception as e: print(f"โš ๏ธ IndexTTS-2 force reload failed: {e}, falling back to full recreation") print(f"๐Ÿ”„ Force reloading {model_type} ({engine}) - removing from cache") self.remove_model(model_key) # CRITICAL: Unload old same-engine model variants BEFORE ComfyUI's free_memory # (prevents loading multiple Qwen3-TTS variants with different attn_implementation, etc.) if model_type == "tts" and engine != "": cached_models = list(self._model_cache.keys()) for existing_key in cached_models: if existing_key.startswith(f"{engine}_tts_") and existing_key != model_key: wrapper = self._model_cache.get(existing_key) if wrapper and wrapper._is_loaded_on_gpu: print(f"๐Ÿ—‘๏ธ Unloading old {engine} model variant to prevent VRAM accumulation") self.remove_model(existing_key) break # Only unload one old variant # Aggressive memory management before loading new model memory_freed = 0 # Track if ComfyUI freed memory (indicates VRAM pressure) if COMFYUI_AVAILABLE and model_management is not None and device != 'cpu': try: # Free up memory aggressively - request 3GB to ensure space for new model if hasattr(model_management, 'free_memory') and callable(getattr(model_management, 'free_memory', None)): if hasattr(model_management, 'get_torch_device'): torch_device = model_management.get_torch_device() # Request 3GB of free VRAM (aggressive cleanup for TTS models) memory_freed = model_management.free_memory(3 * 1024 * 1024 * 1024, torch_device) if memory_freed and memory_freed > 0: print(f"๐Ÿงน Freed {memory_freed // 1024 // 1024}MB VRAM for new {model_type} model") except Exception as e: # Silently ignore memory management errors to avoid spam pass if model_type == "tts" and engine != "" and str(device).startswith('cuda'): self._offload_conflicting_tts_models(engine, device, keep_model_key=model_key) # Create the model # print(f"๐Ÿ”ง Creating new {model_type} model ({engine}) on {device} - fresh instance after cache invalidation") # Higgs Audio now uses deferred CUDA graph initialization to prevent corruption if device.startswith('cuda') and engine == "higgs_audio": # print(f"๐Ÿ“ Creating fresh {engine} model (CUDA graphs deferred until first inference)") try: import gc gc.collect() except Exception as gc_error: print(f"โš ๏ธ Garbage collection failed (safe to ignore): {gc_error}") # Call factory function with config if provided, otherwise use **kwargs if 'config' in factory_kwargs: # New pattern: factory receives ModelLoadConfig object config = factory_kwargs.pop('config') model = model_factory_func(config) else: # Legacy pattern: factory receives **kwargs factory_kwargs['device'] = device model = model_factory_func(**factory_kwargs) # Calculate memory usage memory_size = ComfyUIModelWrapper.calculate_model_memory(model) # Create model info - for stateless wrappers, use a generic engine name to prevent CUDA graph handling actual_engine = engine original_engine = None if hasattr(model, '_wrapped_engine') and engine == "higgs_audio": # This is a stateless wrapper - use generic name to prevent ComfyUI from doing special CUDA handling original_engine = engine # Store original engine name actual_engine = "stateless_tts" model_info = ModelInfo( model=model, model_type=model_type, engine=actual_engine, # Use generic name for stateless wrappers device=device, memory_size=memory_size, load_device=device, original_engine=original_engine # Store original engine before masking ) # Wrap for ComfyUI, passing cache_key so wrapper knows which specific model it is wrapper = ComfyUIModelWrapper(model, model_info, cache_key=model_key) # Cache the wrapper self._model_cache[model_key] = wrapper # Register with ComfyUI using the proper load_models_gpu method if COMFYUI_AVAILABLE and model_management is not None: # Try the safer manual approach first since load_models_gpu seems to have issues try: if hasattr(model_management, 'LoadedModel') and hasattr(model_management, 'current_loaded_models'): loaded_model = model_management.LoadedModel(wrapper) if model is not None: loaded_model.real_model = weakref.ref(model) # Set up the finalizer that ComfyUI expects if hasattr(model_management, 'cleanup_models'): loaded_model.model_finalizer = weakref.finalize(model, model_management.cleanup_models) else: # Create a dummy finalizer that doesn't crash loaded_model.model_finalizer = weakref.finalize(model, lambda: None) else: loaded_model.real_model = weakref.ref(wrapper) loaded_model.model_finalizer = weakref.finalize(wrapper, lambda: None) # Keep a strong reference to our wrapper to prevent garbage collection # This ensures LoadedModel.model property doesn't return None loaded_model._tts_wrapper_ref = wrapper model_management.current_loaded_models.insert(0, loaded_model) # Insert at 0 like ComfyUI does total_models = len(model_management.current_loaded_models) # print(f"๐Ÿ“ฆ {model_type.title()} ready with ComfyUI integration (#{total_models})") else: print(f"โš ๏ธ ComfyUI LoadedModel or current_loaded_models not available") except Exception as e: print(f"โš ๏ธ Failed to register with ComfyUI model management: {e}") return wrapper def get_model(self, model_key: str) -> Optional[ComfyUIModelWrapper]: """Get a cached model by key""" return self._model_cache.get(model_key) def remove_model(self, model_key: str) -> bool: """Remove a model from cache and ComfyUI tracking""" if model_key in self._model_cache: wrapper = self._model_cache[model_key] # With stateless wrapper, Higgs Audio models can now be safely unloaded print(f"๐Ÿ—‘๏ธ Attempting to unload {wrapper.model_info.engine} model (stateless wrapper enabled)") # Measure VRAM before unload vram_before = 0 if wrapper._is_loaded_on_gpu: import torch if torch.cuda.is_available(): vram_before = torch.cuda.memory_allocated() # Normal destruction for all engines self._model_cache.pop(model_key) # Remove from ComfyUI tracking if available if COMFYUI_AVAILABLE and model_management is not None: try: if hasattr(model_management, 'current_loaded_models'): tracked_models = model_management.current_loaded_models matches = [ loaded_model for loaded_model in list(tracked_models) if loaded_model is wrapper or getattr(loaded_model, '_tts_wrapper_ref', None) is wrapper ] for loaded_model in matches: tracked_models.remove(loaded_model) if matches: print(f"๐Ÿ—‘๏ธ Removed model from ComfyUI tracking") except Exception as e: print(f"โš ๏ธ Failed to remove from ComfyUI tracking: {e}") # Permanent removal should release engines that provide a direct # destruction path. Normal ComfyUI offloading still goes through # wrapper.model_unload()/partially_unload() and remains reusable. from .engine_handlers import get_engine_handler handler = get_engine_handler(wrapper.model_info.engine) release = getattr(handler, 'release', None) if callable(release): release(wrapper) else: wrapper.model_unload() # Explicitly delete wrapper to release all references del wrapper # Force garbage collection import gc gc.collect() import torch if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() torch.cuda.empty_cache() # Measure VRAM after unload and report actual freed amount if vram_before > 0: if torch.cuda.is_available(): vram_after = torch.cuda.memory_allocated() freed = vram_before - vram_after print(f"๐Ÿ“Š VRAM: {vram_before // 1024 // 1024}MB โ†’ {vram_after // 1024 // 1024}MB (freed {freed // 1024 // 1024}MB)") return True return False def clear_cache(self, model_type: Optional[str] = None, engine: Optional[str] = None): """Clear cached models with optional filtering""" keys_to_remove = [] for key, wrapper in self._model_cache.items(): should_remove = True if model_type and wrapper.model_info.model_type != model_type: should_remove = False if engine and wrapper.model_info.engine != engine: should_remove = False if should_remove: keys_to_remove.append(key) for key in keys_to_remove: self.remove_model(key) print(f"๐Ÿงน Cleared {len(keys_to_remove)} models from cache") def ensure_device(self, engine_name: str, target_device: str, model_cache_key: str = None) -> bool: """ Ensure a model is on the target device, clearing other models if needed. This is the SINGLE point where device movement happens - all engines should call this instead of doing their own .to(device) calls. Args: engine_name: Name of the engine requesting device movement target_device: Target device ("cuda", "cpu", etc.) model_cache_key: Optional specific model cache key (to distinguish vibevoice-7B from vibevoice-1.5B, etc.) Returns: True if successful, False otherwise """ # Before moving to GPU, ALWAYS move other TTS models to CPU to free VRAM # (Don't delete them - keep in cache for fast reload) if target_device.startswith('cuda'): cleared = self._offload_conflicting_tts_models( engine_name, target_device, keep_model_key=model_cache_key, ) if cleared: print(f"โœ… CLEARING COMPLETE: VRAM freed and ready for {engine_name}") # Find the model for this engine wrapper = None if model_cache_key: # If specific cache key provided, use it directly wrapper = self._model_cache.get(model_cache_key) if not wrapper: print(f"โš ๏ธ ensure_device: Specified model cache key not found: {model_cache_key[:60]}...") return False else: # Fallback: find first model with matching engine (old behavior) for cache_key, cached_wrapper in self._model_cache.items(): if cached_wrapper.model_info.engine == engine_name: wrapper = cached_wrapper break if not wrapper: # Model not in unified cache - engine must be managing it itself # We've already cleared other models, so return False to let engine handle movement return False # Check if already on target device if wrapper.current_device == target_device: return True # Move the model to target device try: # Pass flag to prevent recursion (model_load shouldn't call ensure_device again) wrapper.model_load(target_device, _from_ensure_device=True) return True except Exception as e: print(f"โŒ Failed to move {engine_name} to {target_device}: {e}") return False def get_stats(self) -> Dict[str, Any]: """Get cache statistics""" total_memory = sum(w.loaded_size() for w in self._model_cache.values()) by_type = {} by_engine = {} for wrapper in self._model_cache.values(): model_type = wrapper.model_info.model_type engine = wrapper.model_info.engine by_type[model_type] = by_type.get(model_type, 0) + 1 by_engine[engine] = by_engine.get(engine, 0) + 1 return { 'total_models': len(self._model_cache), 'total_memory_mb': total_memory // 1024 // 1024, 'by_type': by_type, 'by_engine': by_engine, 'comfyui_integration': COMFYUI_AVAILABLE } # Global instance for all TTS models tts_model_manager = ComfyUITTSModelManager()