227 lines
8.3 KiB
Python
227 lines
8.3 KiB
Python
"""
|
|
Engine-specific handlers for model unloading and device management.
|
|
|
|
Different TTS engines have different architectures and device management requirements.
|
|
This module provides specialized handlers for each engine type.
|
|
"""
|
|
|
|
import torch
|
|
import gc
|
|
|
|
|
|
class BaseEngineHandler:
|
|
"""Base handler with default implementation for most engines"""
|
|
|
|
def partially_unload(self, wrapper, device: str, memory_to_free: int) -> int:
|
|
"""
|
|
Partially unload the model to free memory.
|
|
|
|
Args:
|
|
wrapper: ComfyUIModelWrapper instance
|
|
device: Target device (usually 'cpu')
|
|
memory_to_free: Amount of memory to free in bytes
|
|
|
|
Returns:
|
|
Amount of memory actually freed in bytes
|
|
"""
|
|
if not wrapper._is_loaded_on_gpu:
|
|
print(f"⚠️ Skipping unload: model already marked as not on GPU")
|
|
return 0
|
|
|
|
# Get the actual model
|
|
model = wrapper._model_ref() if wrapper._model_ref else None
|
|
if model is None:
|
|
print(f"⚠️ Model reference is None, cannot unload")
|
|
return 0
|
|
|
|
# Move model to CPU
|
|
try:
|
|
if hasattr(model, 'to'):
|
|
model.to(device)
|
|
print(f"🔄 Moved {wrapper.model_info.model_type} model ({wrapper.model_info.engine}) to {device}, freed {wrapper._memory_size // 1024 // 1024}MB")
|
|
|
|
# Update wrapper state
|
|
wrapper.current_device = device
|
|
wrapper._is_loaded_on_gpu = False
|
|
|
|
# Force CUDA cache cleanup
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return wrapper._memory_size
|
|
except Exception as e:
|
|
print(f"⚠️ Error moving model to {device}: {e}")
|
|
return 0
|
|
|
|
def model_unload(self, wrapper, memory_to_free=None, unpatch_weights=True) -> bool:
|
|
"""
|
|
Fully unload the model from GPU memory.
|
|
|
|
Args:
|
|
wrapper: ComfyUIModelWrapper instance
|
|
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: {wrapper.model_info.engine} {wrapper.model_info.model_type}")
|
|
|
|
# Use partially_unload to do the actual work
|
|
freed = self.partially_unload(wrapper, 'cpu', wrapper._memory_size)
|
|
return freed > 0
|
|
|
|
|
|
class StepAudioEditXHandler(BaseEngineHandler):
|
|
"""Specialized handler for Step Audio EditX with bitsandbytes int8 support"""
|
|
|
|
def partially_unload(self, wrapper, device: str, memory_to_free: int) -> int:
|
|
"""
|
|
Unload Step Audio EditX model, handling bitsandbytes int8 models specially.
|
|
|
|
Bitsandbytes int8 models can't use .to() for device movement.
|
|
We need to delete the model and let garbage collection handle it.
|
|
"""
|
|
if not wrapper._is_loaded_on_gpu:
|
|
print(f"⚠️ Skipping unload: model already marked as not on GPU")
|
|
return 0
|
|
|
|
# Get the actual model
|
|
model = wrapper._model_ref() if wrapper._model_ref else None
|
|
if model is None:
|
|
print(f"⚠️ Model reference is None, cannot unload")
|
|
return 0
|
|
|
|
try:
|
|
# For bitsandbytes int8 models, .to() doesn't work
|
|
# Instead, delete the model and force garbage collection
|
|
model_info = f"{wrapper.model_info.model_type} model ({wrapper.model_info.engine})"
|
|
|
|
# Check if this is a bitsandbytes quantized model
|
|
is_quantized = hasattr(model, 'quantization_method') or any(
|
|
'Int8' in str(type(m).__name__) for m in model.modules() if hasattr(model, 'modules')
|
|
)
|
|
|
|
# print(f"DEBUG: Quantization check for {model_info} - is_quantized={is_quantized}, has_quantization_method={hasattr(model, 'quantization_method')}")
|
|
|
|
if is_quantized:
|
|
print(f"🔄 Unloading quantized {model_info} (bitsandbytes int8)...")
|
|
# Delete model reference and let garbage collection handle it
|
|
del model
|
|
wrapper._model_ref = None
|
|
else:
|
|
# Regular model - use standard .to() method
|
|
if hasattr(model, 'to'):
|
|
model.to(device)
|
|
print(f"🔄 Moved {model_info} to {device}")
|
|
|
|
# Update wrapper state
|
|
wrapper.current_device = device
|
|
wrapper._is_loaded_on_gpu = False
|
|
|
|
# Force CUDA cache cleanup
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return wrapper._memory_size
|
|
except Exception as e:
|
|
print(f"⚠️ Error unloading {wrapper.model_info.engine} model: {e}")
|
|
return 0
|
|
|
|
|
|
class HiggsAudioHandler(BaseEngineHandler):
|
|
"""Specialized handler for Higgs Audio engine with CUDA graphs"""
|
|
|
|
def model_unload(self, wrapper, memory_to_free=None, unpatch_weights=True) -> bool:
|
|
"""
|
|
Higgs Audio uses CUDA graphs which are safely cleaned up during device migration.
|
|
CUDA graphs will be recreated automatically on the next generation call.
|
|
"""
|
|
result = super().model_unload(wrapper, memory_to_free, unpatch_weights)
|
|
|
|
if result:
|
|
# Mark model as invalid to force fresh CUDA graph initialization
|
|
wrapper._is_valid_for_reuse = False
|
|
print(f"🔄 Marked Higgs Audio model for cache refresh (CUDA graphs will be recreated)")
|
|
|
|
return result
|
|
|
|
|
|
class DotsTTSHandler(BaseEngineHandler):
|
|
"""Handler for Dots TTS, which needs explicit runtime teardown to free compile caches."""
|
|
|
|
def partially_unload(self, wrapper, device: str, memory_to_free: int) -> int:
|
|
if not wrapper._is_loaded_on_gpu:
|
|
print(f"⚠️ Skipping unload: model already marked as not on GPU")
|
|
return 0
|
|
|
|
model = wrapper._model_ref() if wrapper._model_ref else None
|
|
if model is None:
|
|
print(f"⚠️ Model reference is None, cannot unload")
|
|
return 0
|
|
|
|
try:
|
|
if hasattr(model, "unload_runtime"):
|
|
model.unload_runtime()
|
|
print(
|
|
f"🔄 Unloaded {wrapper.model_info.model_type} model "
|
|
f"({wrapper.model_info.engine}) runtime and freed "
|
|
f"{wrapper._memory_size // 1024 // 1024}MB"
|
|
)
|
|
elif hasattr(model, "to"):
|
|
model.to(device)
|
|
print(
|
|
f"🔄 Moved {wrapper.model_info.model_type} model "
|
|
f"({wrapper.model_info.engine}) to {device}, freed "
|
|
f"{wrapper._memory_size // 1024 // 1024}MB"
|
|
)
|
|
|
|
wrapper.current_device = device
|
|
wrapper._is_loaded_on_gpu = False
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
if hasattr(torch.cuda, "ipc_collect"):
|
|
try:
|
|
torch.cuda.ipc_collect()
|
|
except Exception:
|
|
pass
|
|
gc.collect()
|
|
return wrapper._memory_size
|
|
except Exception as e:
|
|
print(f"⚠️ Error unloading {wrapper.model_info.engine} model: {e}")
|
|
return 0
|
|
|
|
|
|
# Engine registry
|
|
_ENGINE_HANDLERS = {
|
|
'chatterbox': BaseEngineHandler(),
|
|
'chatterbox_official_23lang': BaseEngineHandler(),
|
|
'dots_tts': DotsTTSHandler(),
|
|
'f5tts': BaseEngineHandler(),
|
|
'echo_tts': BaseEngineHandler(),
|
|
'higgs_audio': HiggsAudioHandler(),
|
|
'step_audio_editx': StepAudioEditXHandler(),
|
|
'vibevoice': BaseEngineHandler(),
|
|
'rvc': BaseEngineHandler(),
|
|
}
|
|
|
|
|
|
def get_engine_handler(engine_type: str):
|
|
"""
|
|
Get the appropriate handler for an engine type.
|
|
|
|
Args:
|
|
engine_type: Engine identifier (e.g., 'chatterbox', 'higgs_audio')
|
|
|
|
Returns:
|
|
Engine handler instance
|
|
"""
|
|
handler = _ENGINE_HANDLERS.get(engine_type)
|
|
if handler is None:
|
|
print(f"⚠️ No specialized handler for engine '{engine_type}', using default")
|
|
return BaseEngineHandler()
|
|
return handler
|