diodiogod-TTS-Audio-Suite/utils/models/comfyui_model_wrapper/model_manager.py

"""
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 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")

                # Also try manual cleanup of our own TTS model cache
                # Only if ComfyUI actually had to free memory (VRAM was tight)
                if model_type == "tts" and engine != "" and memory_freed > 0:
                    # Get current cache stats
                    cached_models = list(self._model_cache.keys())
                    models_to_offload = []

                    for cache_key in cached_models:
                        wrapper = self._model_cache[cache_key]

                        # Move models to CPU to free VRAM if:
                        # 1. Different engine, OR
                        # 2. Same engine but different model variant (different cache key)
                        # (Don't delete - keep in cache for fast reload)
                        should_offload = (
                            wrapper.model_info.model_type == "tts" and
                            wrapper._is_loaded_on_gpu and
                            (wrapper.model_info.engine != engine or cache_key != model_key)
                        )
                        if should_offload:
                            models_to_offload.append((cache_key, wrapper))

                    if models_to_offload:
                        print(f"🗑️ Moving {len(models_to_offload)} TTS models to CPU to free VRAM")
                        for 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: {e}")
                                # Keep in cache anyway - might recover later

                        # Force CUDA to synchronize and free memory before loading next model
                        import torch
                        if torch.cuda.is_available():
                            torch.cuda.synchronize()
                            torch.cuda.empty_cache()
                            import time
                            time.sleep(0.1)
                            
            except Exception as e:
                # Silently ignore memory management errors to avoid spam
                pass
        
        # 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"
            print(f"🔒 Treating {engine} stateless wrapper as generic TTS model to avoid CUDA graph interference")

        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:
                # Manually add to current_loaded_models using LoadedModel (ComfyUI's internal approach)
                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)
                    
                    # 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'):
                        if wrapper in model_management.current_loaded_models:
                            model_management.current_loaded_models.remove(wrapper)
                            print(f"🗑️ Removed model from ComfyUI tracking")
                except Exception as e:
                    print(f"⚠️ Failed to remove from ComfyUI tracking: {e}")

            # Unload from GPU
            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'):
            cached_models = list(self._model_cache.keys())
            models_to_offload = []

            for cache_key in cached_models:
                cached_wrapper = self._model_cache.get(cache_key)
                # Clear models from different engines OR same engine but different cache key
                # (e.g., vibevoice-7B on GPU should be cleared when loading vibevoice-1.5B)
                should_offload = (
                    cached_wrapper and
                    cached_wrapper._is_loaded_on_gpu and
                    (cached_wrapper.model_info.engine != engine_name or
                     (model_cache_key is not None and cache_key != model_cache_key))
                )
                if should_offload:
                    models_to_offload.append((cache_key, cached_wrapper))

            if models_to_offload:
                print(f"🗑️ BEFORE CLEARING: Moving {len(models_to_offload)} TTS models to CPU to make room for {engine_name}")
                for key, wrapper in models_to_offload:
                    # Move to CPU instead of deleting - keeps model in cache for fast reload
                    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(key)

                # CRITICAL: Force CUDA to synchronize and actually free the memory before loading next model
                # Without this, the next model's .to(cuda) may fail with OOM because CUDA hasn't freed yet
                print(f"⏳ AFTER CLEARING: Synchronizing CUDA to ensure memory is freed...")
                import torch
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                    torch.cuda.empty_cache()
                    import time
                    time.sleep(0.1)  # Brief pause to ensure GPU memory is actually freed
                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()