diodiogod
d01e5cb807
- Split 970-line monolith into focused modules following 500-600 line policy
- Created utils/models/comfyui_model_wrapper/ directory structure:
- base_wrapper.py (302 lines): Core ComfyUIModelWrapper class
- model_manager.py (318 lines): ComfyUITTSModelManager class
- cache_utils.py (28 lines): Cache invalidation utilities
- engine_handlers/: Engine-specific memory management
- higgs_audio_handler.py: CUDA graph management and unload blocking
- vibevoice_handler.py: RAM cleanup and complete deletion logic
- generic_handler.py: Standard CPU migration for other engines
- Fixed engine routing bug where stateless_tts wrappers weren't using HiggsAudioHandler
- Maintained full backward compatibility with existing imports
- All engine-specific behaviors preserved (VibeVoice deletion, Higgs CUDA graphs)
- Updated PROJECT_INDEX.md to reflect modular structure
970 lines
47 KiB
Plaintext
970 lines
47 KiB
Plaintext
"""
|
|
ComfyUI Model Wrapper for TTS Audio Suite
|
|
|
|
This module provides wrappers that make TTS models compatible with ComfyUI's
|
|
model management system, enabling automatic memory management, "Clear VRAM"
|
|
button functionality, and proper integration with ComfyUI's model lifecycle.
|
|
"""
|
|
|
|
import torch
|
|
import weakref
|
|
import gc
|
|
import time
|
|
from typing import Optional, Any, Dict, Union
|
|
from dataclasses import dataclass
|
|
|
|
# Global cache invalidation flag to force recreation of all engine instances
|
|
# When models are unloaded, this timestamp is updated to invalidate all node caches
|
|
_global_cache_invalidation_flag = 0.0
|
|
|
|
|
|
def is_engine_cache_valid(cache_timestamp: float) -> bool:
|
|
"""
|
|
Check if an engine cache is still valid based on global invalidation flag.
|
|
|
|
Args:
|
|
cache_timestamp: When the cache entry was created
|
|
|
|
Returns:
|
|
True if cache is still valid, False if it should be invalidated
|
|
"""
|
|
return cache_timestamp > _global_cache_invalidation_flag
|
|
|
|
# Import ComfyUI's model management if available
|
|
try:
|
|
import comfy.model_management as model_management
|
|
COMFYUI_AVAILABLE = True
|
|
except ImportError:
|
|
# Fallback if ComfyUI not available
|
|
COMFYUI_AVAILABLE = False
|
|
model_management = None
|
|
|
|
|
|
@dataclass
|
|
class ModelInfo:
|
|
"""Information about a model for memory management"""
|
|
model: Any
|
|
model_type: str # "tts", "vc", "audio_separation", "hubert", etc.
|
|
engine: str # "chatterbox", "f5tts", "higgs_audio", "rvc", etc.
|
|
device: str
|
|
memory_size: int # in bytes
|
|
load_device: str
|
|
|
|
|
|
class ComfyUIModelWrapper:
|
|
"""
|
|
Wrapper that makes TTS models compatible with ComfyUI's model management system.
|
|
|
|
This allows TTS models to be automatically unloaded when VRAM is low,
|
|
work with "Clear VRAM" buttons, and integrate properly with ComfyUI's ecosystem.
|
|
"""
|
|
|
|
def __init__(self, model: Any, model_info: ModelInfo):
|
|
"""
|
|
Initialize the wrapper.
|
|
|
|
Args:
|
|
model: The actual model instance (ChatterBox, F5-TTS, etc.)
|
|
model_info: Metadata about the model
|
|
"""
|
|
self.model = model
|
|
self.model_info = model_info
|
|
self.load_device = model_info.load_device
|
|
self.current_device = model_info.device
|
|
self._memory_size = model_info.memory_size
|
|
|
|
# ComfyUI compatibility attributes
|
|
# Convert device to torch.device object for ComfyUI compatibility
|
|
device_name = model_info.device
|
|
if device_name == "auto":
|
|
device_name = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# ComfyUI expects torch.device objects, not strings
|
|
if isinstance(device_name, str):
|
|
if device_name == "cuda":
|
|
self.device = torch.device("cuda", torch.cuda.current_device() if torch.cuda.is_available() else 0)
|
|
else:
|
|
self.device = torch.device(device_name)
|
|
else:
|
|
self.device = device_name
|
|
self.dtype = getattr(model, 'dtype', torch.float32)
|
|
self.offload_device = 'cpu' # TTS models offload to CPU
|
|
|
|
# ComfyUI compatibility attributes for diffusion models (TTS models don't need them)
|
|
self.model_patches_models = [] # Empty list for TTS models
|
|
self.parent = None # TTS models don't have parent models
|
|
|
|
# Additional ComfyUI LoadedModel compatibility attributes
|
|
# Use the same torch.device object for load_device
|
|
self.load_device = self.device
|
|
self.currently_used = True
|
|
self.model_finalizer = None # Will be set by LoadedModel
|
|
self._patcher_finalizer = None
|
|
|
|
# ComfyUI model patcher attributes (required for load_models_gpu)
|
|
self.model_patches_to = {} # Patch mapping for diffusion models (empty for TTS)
|
|
self.model_options = {} # Model loading options
|
|
self.model_keys = set() # Model state dict keys
|
|
|
|
# Track if model is currently loaded on GPU
|
|
self._is_loaded_on_gpu = self.current_device not in ['cpu', 'offload']
|
|
|
|
# Track if model is valid for reuse (false if corrupted by CPU offloading)
|
|
self._is_valid_for_reuse = True
|
|
|
|
# Keep weak reference to avoid circular references
|
|
self._model_ref = weakref.ref(model) if model is not None else None
|
|
|
|
def loaded_size(self) -> int:
|
|
"""Return the memory size of the model in bytes"""
|
|
size = self._memory_size if self._is_loaded_on_gpu else 0
|
|
return size
|
|
|
|
def model_size(self) -> int:
|
|
"""Return the total model size in bytes"""
|
|
return self._memory_size
|
|
|
|
def model_offloaded_memory(self) -> int:
|
|
"""Return the amount of memory that would be freed if offloaded"""
|
|
return self.model_size() - self.loaded_size()
|
|
|
|
def current_loaded_device(self) -> str:
|
|
"""Return the current device the model is loaded on"""
|
|
return self.current_device
|
|
|
|
def partially_unload(self, device: str, memory_to_free: int) -> int:
|
|
"""
|
|
Partially unload the model to free memory.
|
|
|
|
For TTS models, this typically means moving to CPU or offloading.
|
|
For VibeVoice, we delete completely to prevent system RAM accumulation.
|
|
|
|
Args:
|
|
device: Target device to move to (usually 'cpu')
|
|
memory_to_free: Amount of memory to free in bytes
|
|
|
|
Returns:
|
|
Amount of memory actually freed in bytes
|
|
"""
|
|
if not self._is_loaded_on_gpu:
|
|
return 0
|
|
|
|
# VibeVoice special handling: Smart CPU migration with RAM cleanup
|
|
if self.model_info.engine == "vibevoice" and device == 'cpu':
|
|
print(f"🔄 VibeVoice: Smart CPU migration with RAM cleanup to prevent accumulation")
|
|
|
|
# Before moving to CPU, clean up any existing VibeVoice models in system RAM
|
|
# Simple approach: clear other VibeVoice models that aren't on GPU
|
|
try:
|
|
models_cleared = 0
|
|
from utils.models.comfyui_model_wrapper import tts_model_manager
|
|
cache_keys_to_clear = []
|
|
|
|
# Find VibeVoice models that are in CPU/RAM (not GPU loaded)
|
|
for cache_key, wrapper in tts_model_manager._model_cache.items():
|
|
if (wrapper.model_info.engine == "vibevoice" and
|
|
not wrapper._is_loaded_on_gpu and # Model is in RAM/CPU
|
|
wrapper != self): # Don't clear ourselves
|
|
cache_keys_to_clear.append(cache_key)
|
|
models_cleared += 1
|
|
|
|
# Clear the old VibeVoice models from RAM
|
|
for key in cache_keys_to_clear:
|
|
try:
|
|
tts_model_manager.remove_model(key)
|
|
print(f"🗑️ Removed VibeVoice model from RAM: {key[:16]}...")
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to remove {key[:16]}: {e}")
|
|
|
|
if models_cleared > 0:
|
|
print(f"🧹 RAM cleanup: removed {models_cleared} old VibeVoice models from system memory")
|
|
import gc
|
|
gc.collect()
|
|
else:
|
|
print(f"🔍 No old VibeVoice models found in RAM")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ RAM cleanup error: {e}")
|
|
|
|
# Now proceed with normal CPU migration
|
|
print(f"📥 Moving VibeVoice to CPU (RAM cleanup completed)")
|
|
|
|
model = self._model_ref() if self._model_ref else None
|
|
if model is None:
|
|
return 0
|
|
|
|
freed_memory = 0
|
|
|
|
try:
|
|
# Move model to CPU if it has a .to() method
|
|
if hasattr(model, 'to'):
|
|
try:
|
|
# CRITICAL: Clear CUDA graphs before moving to CPU (prevents corruption)
|
|
self._clear_cuda_graphs(model)
|
|
|
|
model.to('cpu')
|
|
freed_memory = self._memory_size
|
|
self.current_device = 'cpu'
|
|
self._is_loaded_on_gpu = False
|
|
print(f"🔄 Moved {self.model_info.model_type} model ({self.model_info.engine}) to CPU, freed {freed_memory // 1024 // 1024}MB")
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to move {self.model_info.model_type} model to CPU: {e}")
|
|
# Still mark as unloaded if the model reported an error moving to CPU
|
|
self.current_device = 'cpu'
|
|
self._is_loaded_on_gpu = False
|
|
freed_memory = self._memory_size
|
|
|
|
# Handle nested models (like ChatterBox with multiple components)
|
|
elif hasattr(model, '__dict__'):
|
|
# CRITICAL: Clear CUDA graphs before moving to CPU (prevents corruption)
|
|
self._clear_cuda_graphs(model)
|
|
|
|
for attr_name, attr_value in model.__dict__.items():
|
|
if hasattr(attr_value, 'to') and hasattr(attr_value, 'parameters'):
|
|
try:
|
|
attr_value.to('cpu')
|
|
freed_memory += self._estimate_model_memory(attr_value)
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to move {attr_name} to CPU: {e}")
|
|
pass
|
|
|
|
if freed_memory > 0:
|
|
self.current_device = 'cpu'
|
|
self._is_loaded_on_gpu = False
|
|
print(f"🔄 Moved {self.model_info.model_type} model components ({self.model_info.engine}) to CPU, freed {freed_memory // 1024 // 1024}MB")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to partially unload {self.model_info.model_type} model: {e}")
|
|
|
|
# Force garbage collection after unloading
|
|
if freed_memory > 0:
|
|
gc.collect()
|
|
# Be more careful with CUDA cache clearing to avoid interfering with CUDA graphs
|
|
if torch.cuda.is_available():
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f"⚠️ CUDA cache clear warning (safe to ignore): {e}")
|
|
|
|
return freed_memory
|
|
|
|
def model_unload(self, memory_to_free: Optional[int] = None, unpatch_weights: bool = True) -> bool:
|
|
"""
|
|
Fully unload the model from GPU memory.
|
|
|
|
Args:
|
|
memory_to_free: Amount of memory to free (ignored for full unload)
|
|
unpatch_weights: Whether to unpatch weights (TTS models don't use this)
|
|
|
|
Returns:
|
|
True if model was unloaded, False otherwise
|
|
"""
|
|
print(f"🔄 TTS Model unload requested: {self.model_info.engine} {self.model_info.model_type}")
|
|
|
|
# Check if this is a Higgs Audio model with CUDA Graphs enabled
|
|
if self.model_info.engine == "higgs_audio":
|
|
cuda_graphs_enabled = getattr(self.model, '_cuda_graphs_enabled', True)
|
|
if cuda_graphs_enabled:
|
|
print(f"⛔ CUDA Graph Mode: Unloading disabled to prevent crashes")
|
|
print(f" Model uses CUDA Graph optimization - cannot be safely unloaded")
|
|
print(f" To enable memory unloading, disable CUDA Graphs in engine settings")
|
|
print(f" Or restart ComfyUI to fully free memory")
|
|
return False # Refuse to unload
|
|
|
|
if memory_to_free is not None and memory_to_free < self.loaded_size():
|
|
# Try partial unload first
|
|
freed = self.partially_unload('cpu', memory_to_free)
|
|
success = freed >= memory_to_free
|
|
print(f"{'✅' if success else '❌'} Partial unload: freed {freed // 1024 // 1024}MB (requested {memory_to_free // 1024 // 1024}MB)")
|
|
return success
|
|
|
|
# Full unload - for VibeVoice, delete completely instead of moving to CPU
|
|
if self.model_info.engine == "vibevoice":
|
|
# VibeVoice: Full deletion to prevent system RAM accumulation
|
|
freed = 0
|
|
model_location = "unknown"
|
|
|
|
if hasattr(self, 'model') and self.model is not None:
|
|
try:
|
|
# Detect if model is on GPU or CPU
|
|
if hasattr(self.model, 'device'):
|
|
model_location = str(self.model.device)
|
|
elif self._is_loaded_on_gpu:
|
|
model_location = "GPU"
|
|
else:
|
|
model_location = "CPU"
|
|
|
|
print(f"🔍 VibeVoice deletion: Model currently on {model_location}")
|
|
|
|
# Clear VibeVoice class-level cache first
|
|
if hasattr(self.model, '__class__'):
|
|
model_class = self.model.__class__
|
|
if hasattr(model_class, '_shared_model'):
|
|
model_class._shared_model = None
|
|
model_class._shared_processor = None
|
|
model_class._shared_config = None
|
|
model_class._shared_model_name = None
|
|
print(f"🗑️ Cleared VibeVoice class-level cache")
|
|
|
|
# Estimate memory before deletion (same size regardless of location)
|
|
freed = self._memory_size
|
|
|
|
# Delete the model completely (works for both GPU and CPU)
|
|
del self.model
|
|
self.model = None
|
|
|
|
# Force garbage collection to actually free memory
|
|
import gc
|
|
gc.collect()
|
|
|
|
# Clear CUDA cache if model was on GPU - AGGRESSIVE CLEANUP
|
|
if model_location.lower() in ['gpu', 'cuda'] or 'cuda' in model_location.lower():
|
|
if hasattr(torch, 'cuda') and torch.cuda.is_available():
|
|
# Force multiple cleanup passes
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize() # Wait for all operations to complete
|
|
torch.cuda.empty_cache() # Second pass after sync
|
|
|
|
# Force garbage collection multiple times
|
|
for _ in range(3): # Multiple GC passes
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# NUCLEAR OPTION: Force CUDA device reset
|
|
try:
|
|
print(f"⚠️ NUCLEAR: Forcing CUDA device reset to clear stubborn VibeVoice memory")
|
|
torch.cuda.reset_peak_memory_stats()
|
|
torch.cuda.empty_cache()
|
|
|
|
# Try ComfyUI model management
|
|
try:
|
|
from comfy import model_management
|
|
if hasattr(model_management, 'free_memory'):
|
|
model_management.free_memory(8 * 1024 * 1024 * 1024, torch.cuda.current_device()) # Request 8GB
|
|
print(f"🧹 ComfyUI freed memory")
|
|
except:
|
|
pass
|
|
|
|
# Final synchronization
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ CUDA reset failed: {e}")
|
|
|
|
print(f"🧹 NUCLEAR CUDA cleanup completed")
|
|
|
|
print(f"🗑️ VibeVoice: Model deleted completely from {model_location}")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ VibeVoice deletion error: {e}")
|
|
|
|
success = freed > 0
|
|
print(f"{'✅' if success else '❌'} VibeVoice full deletion: freed {freed // 1024 // 1024}MB from {model_location}")
|
|
return success
|
|
else:
|
|
# Other engines: use standard CPU migration
|
|
freed = self.partially_unload('cpu', self._memory_size)
|
|
success = freed > 0
|
|
print(f"{'✅' if success else '❌'} Full unload: freed {freed // 1024 // 1024}MB")
|
|
return success
|
|
|
|
def _vibevoice_complete_deletion(self) -> bool:
|
|
"""Complete deletion of VibeVoice model to prevent system RAM accumulation"""
|
|
try:
|
|
model = self._model_ref() if self._model_ref else None
|
|
if model is None:
|
|
return False
|
|
|
|
# Clear VibeVoice class-level cache first
|
|
if hasattr(model, '__class__'):
|
|
model_class = model.__class__
|
|
if hasattr(model_class, '_shared_model'):
|
|
model_class._shared_model = None
|
|
model_class._shared_processor = None
|
|
model_class._shared_config = None
|
|
model_class._shared_model_name = None
|
|
|
|
# Delete the model completely
|
|
del model
|
|
self._model_ref = None
|
|
|
|
# Force garbage collection
|
|
import gc
|
|
gc.collect()
|
|
|
|
# Clear CUDA cache
|
|
if hasattr(torch, 'cuda') and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
self._is_loaded_on_gpu = False
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ VibeVoice complete deletion error: {e}")
|
|
return False
|
|
|
|
def _clear_cuda_graphs(self, model):
|
|
"""Clear CUDA graphs if the model supports it (prevents corruption when moving to CPU)"""
|
|
try:
|
|
# Check if this is a Higgs Audio model with CUDA graphs
|
|
if self.model_info.engine == "higgs_audio":
|
|
print(f"🔍 Checking for CUDA graphs in {self.model_info.engine} model...")
|
|
print(f"🔍 Found Higgs Audio model, searching for decode_graph_runners...")
|
|
|
|
# The CUDA graphs are nested deeper in the Higgs Audio model structure
|
|
# Try to find them through various paths
|
|
cuda_model = None
|
|
|
|
# Path 1: Direct access
|
|
if hasattr(model, 'decode_graph_runners'):
|
|
cuda_model = model
|
|
print(f"🔍 Found decode_graph_runners at top level")
|
|
|
|
# Path 2: Through engine attribute
|
|
elif hasattr(model, 'engine') and hasattr(model.engine, 'model') and hasattr(model.engine.model, 'decode_graph_runners'):
|
|
cuda_model = model.engine.model
|
|
print(f"🔍 Found decode_graph_runners in model.engine.model")
|
|
|
|
# Path 3: Through model attribute
|
|
elif hasattr(model, 'model') and hasattr(model.model, 'decode_graph_runners'):
|
|
cuda_model = model.model
|
|
print(f"🔍 Found decode_graph_runners in model.model")
|
|
|
|
# Path 4: Search through all attributes recursively
|
|
else:
|
|
print(f"🔍 Searching recursively for decode_graph_runners...")
|
|
def find_cuda_model(obj, depth=0, max_depth=3):
|
|
if depth > max_depth:
|
|
return None
|
|
if hasattr(obj, 'decode_graph_runners'):
|
|
return obj
|
|
if hasattr(obj, '__dict__'):
|
|
for attr_name, attr_value in obj.__dict__.items():
|
|
if not attr_name.startswith('_') and attr_value is not None:
|
|
result = find_cuda_model(attr_value, depth + 1, max_depth)
|
|
if result:
|
|
print(f"🔍 Found decode_graph_runners in {attr_name} (depth {depth + 1})")
|
|
return result
|
|
return None
|
|
|
|
cuda_model = find_cuda_model(model)
|
|
|
|
if cuda_model:
|
|
# Check for CUDA graphs and try to safely release them
|
|
graph_count = sum(len(runners) for runners in cuda_model.decode_graph_runners.values())
|
|
if graph_count > 0:
|
|
print(f"🔍 Found {graph_count} CUDA graphs - attempting safe release")
|
|
try:
|
|
# Try to properly end/reset the CUDA graphs before clearing
|
|
# This should release the captured allocations properly
|
|
for key, runners in cuda_model.decode_graph_runners.items():
|
|
print(f" 🔧 Releasing {len(runners)} graphs for {key}")
|
|
for i, runner in enumerate(runners):
|
|
if hasattr(runner, 'graph') and runner.graph is not None:
|
|
# Try to reset/end the graph properly
|
|
try:
|
|
# Reset the graph state
|
|
if hasattr(runner.graph, 'reset'):
|
|
runner.graph.reset()
|
|
elif hasattr(runner, 'reset'):
|
|
runner.reset()
|
|
print(f" ✅ Released graph {i+1}/{len(runners)}")
|
|
except Exception as e:
|
|
print(f" ⚠️ Failed to reset graph {i+1}: {e}")
|
|
|
|
# Now clear the runners
|
|
runners.clear()
|
|
|
|
print(f"🧹 Attempted to release {graph_count} CUDA graphs safely")
|
|
|
|
# Force CUDA synchronization to ensure graphs are properly released
|
|
if torch.cuda.is_available():
|
|
torch.cuda.synchronize()
|
|
print(f"🔄 CUDA synchronized after graph release")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to release CUDA graphs: {e}, proceeding with standard unload")
|
|
else:
|
|
print(f"📝 No CUDA graphs found in {self.model_info.engine} model")
|
|
else:
|
|
print(f"⚠️ Could not locate decode_graph_runners in {self.model_info.engine} model structure")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to clear CUDA graphs: {e}")
|
|
|
|
def model_load(self, device: Optional[str] = None) -> None:
|
|
"""
|
|
Load the model back to GPU.
|
|
|
|
Args:
|
|
device: Device to load to (defaults to original load_device)
|
|
"""
|
|
if self._is_loaded_on_gpu:
|
|
return
|
|
|
|
target_device = device or self.load_device
|
|
model = self._model_ref() if self._model_ref else None
|
|
|
|
if model is None:
|
|
return
|
|
|
|
try:
|
|
# Move model back to GPU (comprehensive approach)
|
|
if hasattr(model, 'to'):
|
|
model.to(target_device)
|
|
print(f"🔄 Moved main {self.model_info.model_type} model ({self.model_info.engine}) to {target_device}")
|
|
|
|
# CRITICAL: Recursively move ALL nested components to ensure device consistency
|
|
self._move_all_components_to_device(model, target_device, depth=0)
|
|
|
|
self.current_device = target_device
|
|
self._is_loaded_on_gpu = True
|
|
print(f"✅ Fully moved {self.model_info.model_type} model components ({self.model_info.engine}) back to {target_device}")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Error moving model to {target_device}: {e}")
|
|
|
|
def _move_all_components_to_device(self, obj, target_device: str, depth: int = 0, max_depth: int = 4):
|
|
"""
|
|
Recursively move all PyTorch components to target device.
|
|
This ensures no CPU/GPU device mismatches after model reload.
|
|
"""
|
|
if depth > max_depth:
|
|
return
|
|
|
|
if obj is None:
|
|
return
|
|
|
|
# Move PyTorch modules
|
|
if hasattr(obj, 'to') and hasattr(obj, 'parameters') and callable(getattr(obj, 'to')):
|
|
try:
|
|
obj.to(target_device)
|
|
if depth == 0:
|
|
print(f" 📦 Moved {type(obj).__name__} to {target_device}")
|
|
except Exception as e:
|
|
if depth == 0:
|
|
print(f" ⚠️ Failed to move {type(obj).__name__}: {e}")
|
|
|
|
# Recurse through object attributes
|
|
if hasattr(obj, '__dict__'):
|
|
for attr_name, attr_value in obj.__dict__.items():
|
|
if not attr_name.startswith('_') and attr_value is not None:
|
|
# Skip certain problematic attributes
|
|
if attr_name in ['_modules', '_parameters', '_buffers']:
|
|
continue
|
|
try:
|
|
self._move_all_components_to_device(attr_value, target_device, depth + 1, max_depth)
|
|
except Exception:
|
|
pass
|
|
|
|
def is_clone(self, other) -> bool:
|
|
"""Check if this model is a clone of another model"""
|
|
if not isinstance(other, ComfyUIModelWrapper):
|
|
return False
|
|
return (self.model_info.model_type == other.model_info.model_type and
|
|
self.model_info.engine == other.model_info.engine)
|
|
|
|
def detach(self, unpatch_all: bool = False) -> None:
|
|
"""Detach the model - actually unload from GPU to CPU and invalidate cache"""
|
|
print(f"🔧 TTS Model detach called: {self.model_info.engine} {self.model_info.model_type} (unpatch_all={unpatch_all})")
|
|
|
|
# Actually unload the model from GPU
|
|
freed = self.partially_unload('cpu', self._memory_size)
|
|
if freed > 0:
|
|
print(f"✅ TTS Model detached: freed {freed // 1024 // 1024}MB VRAM")
|
|
else:
|
|
print(f"⚠️ TTS Model detach: no memory freed (model may already be on CPU)")
|
|
|
|
# CRITICAL: Mark model as invalid to prevent reuse of corrupted state
|
|
# Models with CUDA graphs cannot be safely reused after CPU offloading
|
|
self._is_valid_for_reuse = False
|
|
print(f"🚫 Marked {self.model_info.engine} model as invalid for reuse (CUDA state corrupted)")
|
|
|
|
# CRITICAL: Clear node-level engine caches to prevent reuse of corrupted engines
|
|
# This is essential because TTS nodes have their own caching separate from ComfyUI wrapper cache
|
|
self._clear_node_engine_caches()
|
|
|
|
def _clear_node_engine_caches(self):
|
|
"""Clear engine caches in TTS nodes to prevent reuse of corrupted engines"""
|
|
try:
|
|
# Set global flag to invalidate all caches
|
|
# This will be checked by nodes when they try to reuse cached engines
|
|
global _global_cache_invalidation_flag
|
|
_global_cache_invalidation_flag = time.time()
|
|
print(f"🗑️ Set global cache invalidation flag to force engine recreation")
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to clear node engine caches: {e}")
|
|
|
|
def is_clone(self, other) -> bool:
|
|
"""Check if this model is a clone of another model"""
|
|
return False # TTS models don't support cloning
|
|
|
|
def partially_load(self, device, extra_memory, force_patch_weights=False):
|
|
"""
|
|
Partially load model to device (ComfyUI compatibility method)
|
|
|
|
Args:
|
|
device: Target device
|
|
extra_memory: Extra memory available
|
|
force_patch_weights: Whether to force patch weights
|
|
|
|
Returns:
|
|
Amount of memory used
|
|
"""
|
|
print(f"🔄 TTS Model partially_load requested: {self.model_info.engine} {self.model_info.model_type} to {device}")
|
|
|
|
# For TTS models, we either fully load or fully unload
|
|
if device != 'cpu' and not self._is_loaded_on_gpu:
|
|
self.model_load(device)
|
|
return self._memory_size
|
|
|
|
return 0 # No additional memory used
|
|
|
|
@staticmethod
|
|
def _estimate_model_memory(model) -> int:
|
|
"""Estimate memory usage of a PyTorch model"""
|
|
if not hasattr(model, 'parameters'):
|
|
return 0
|
|
|
|
total_size = 0
|
|
for param in model.parameters():
|
|
total_size += param.nelement() * param.element_size()
|
|
return total_size
|
|
|
|
@staticmethod
|
|
def calculate_model_memory(model: Any) -> int:
|
|
"""Calculate total memory usage of a model in bytes"""
|
|
if hasattr(model, 'parameters'):
|
|
# PyTorch model
|
|
return ComfyUIModelWrapper._estimate_model_memory(model)
|
|
elif hasattr(model, '__dict__'):
|
|
# Complex model with multiple components
|
|
total_size = 0
|
|
for attr_name, attr_value in model.__dict__.items():
|
|
if hasattr(attr_value, 'parameters'):
|
|
total_size += ComfyUIModelWrapper._estimate_model_memory(attr_value)
|
|
elif attr_name == 'f5tts_model' and hasattr(attr_value, '__dict__'):
|
|
# Special handling for ChatterBoxF5TTS wrapper
|
|
for sub_attr_name, sub_attr_value in attr_value.__dict__.items():
|
|
if hasattr(sub_attr_value, 'parameters'):
|
|
total_size += ComfyUIModelWrapper._estimate_model_memory(sub_attr_value)
|
|
return total_size if total_size > 0 else 1024 * 1024 * 1024 # 1GB default if nothing found
|
|
else:
|
|
# Estimate based on common model sizes
|
|
return 1024 * 1024 * 1024 # Default 1GB estimate
|
|
|
|
def __repr__(self):
|
|
return f"ComfyUIModelWrapper({self.model_info.model_type}:{self.model_info.engine}, {self._memory_size // 1024 // 1024}MB, device={self.current_device})"
|
|
|
|
|
|
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")
|
|
|
|
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})")
|
|
# 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")
|
|
|
|
print(f"🔄 Force reloading {model_type} ({engine}) - removing from cache")
|
|
self.remove_model(model_key)
|
|
|
|
# Aggressive memory management before loading new model
|
|
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
|
|
# Clear models from other engines to make room
|
|
if model_type == "tts" and engine != "":
|
|
# Get current cache stats
|
|
cached_models = list(self._model_cache.keys())
|
|
models_to_clear = []
|
|
|
|
for cache_key in cached_models:
|
|
wrapper = self._model_cache[cache_key]
|
|
should_clear = False
|
|
|
|
# Clear models from different engines to free VRAM
|
|
if wrapper.model_info.engine != engine and wrapper.model_info.model_type == "tts":
|
|
should_clear = True
|
|
|
|
# VibeVoice-specific: Use CPU migration instead of direct clearing
|
|
# This allows model reuse while preventing RAM accumulation
|
|
elif engine == "vibevoice" and wrapper.model_info.engine == "vibevoice":
|
|
if cache_key != model_key: # Don't clear the model we're about to load
|
|
print(f"🔄 VibeVoice: Moving existing model to CPU instead of clearing (allows reuse)")
|
|
# Force CPU migration instead of deletion
|
|
try:
|
|
wrapper.partially_unload('cpu', wrapper._memory_size)
|
|
except:
|
|
should_clear = True # Fallback to clearing if CPU migration fails
|
|
|
|
if should_clear:
|
|
models_to_clear.append(cache_key)
|
|
|
|
if models_to_clear:
|
|
print(f"🗑️ Clearing {len(models_to_clear)} TTS models to free VRAM")
|
|
for key in models_to_clear:
|
|
self.remove_model(key)
|
|
|
|
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)")
|
|
import gc
|
|
gc.collect()
|
|
|
|
# Ensure device parameter is available to factory function
|
|
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
|
|
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
|
|
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
|
|
)
|
|
|
|
# Wrap for ComfyUI
|
|
wrapper = ComfyUIModelWrapper(model, model_info)
|
|
|
|
# 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)")
|
|
|
|
# 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()
|
|
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 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() |