diodiogod
bb9a64f564
- Consolidate character voice info into block generation messages - Fix tokenizer path structure to match actual storage location (vibevoice/ instead of VibeVoice/tokenizer/) - Use consistent path for tokenizer downloads so they're found on subsequent runs - Comment out non-error device_map debug messages - Remove duplicate SINGLE SEGMENT text printing - Fix logged_characters to only log once per character (not per priority check)
812 lines
34 KiB
Python
812 lines
34 KiB
Python
"""
|
|
Unified Model Interface for TTS Audio Suite
|
|
|
|
This module provides a standardized interface for all TTS engines to use
|
|
ComfyUI's model management system. All engines (ChatterBox, F5-TTS, Higgs Audio,
|
|
RVC, etc.) should use this interface instead of direct caching.
|
|
"""
|
|
|
|
import torch
|
|
from typing import Any, Dict, Optional, Callable, Union
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from utils.models.comfyui_model_wrapper import tts_model_manager, ModelInfo
|
|
|
|
|
|
@dataclass
|
|
class UnifiedModelConfig:
|
|
"""Configuration for unified model loading"""
|
|
engine_name: str # "chatterbox", "f5tts", "higgs_audio", "rvc"
|
|
model_type: str # "tts", "vc", "tokenizer", "hubert", "separator"
|
|
model_name: str # Model identifier
|
|
device: str # Target device
|
|
language: Optional[str] = "English" # For multilingual models
|
|
model_path: Optional[str] = None # Local path if available
|
|
repo_id: Optional[str] = None # HuggingFace repo ID
|
|
additional_params: Optional[Dict[str, Any]] = None # Engine-specific params
|
|
|
|
|
|
class UnifiedModelInterface:
|
|
"""
|
|
Unified interface for all TTS engine model loading.
|
|
|
|
This replaces all engine-specific caching systems with a single
|
|
ComfyUI-integrated approach.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the unified interface"""
|
|
self._model_factories: Dict[str, Callable] = {}
|
|
self._pytorch_warning_shown = False
|
|
|
|
def register_model_factory(self,
|
|
engine_name: str,
|
|
model_type: str,
|
|
factory_func: Callable) -> None:
|
|
"""
|
|
Register a model factory function for an engine.
|
|
|
|
Args:
|
|
engine_name: Name of the engine ("chatterbox", "f5tts", etc.)
|
|
model_type: Type of model ("tts", "vc", "tokenizer", etc.)
|
|
factory_func: Function that creates the model
|
|
"""
|
|
key = f"{engine_name}_{model_type}"
|
|
self._model_factories[key] = factory_func
|
|
|
|
def _check_pytorch_consistency(self) -> None:
|
|
"""Check for PyTorch/TorchAudio/TorchVision consistency and warn if mixed CUDA/CPU"""
|
|
if self._pytorch_warning_shown:
|
|
return
|
|
|
|
try:
|
|
import torch
|
|
import torchaudio
|
|
import torchvision
|
|
|
|
# Check PyTorch version info
|
|
torch_version = getattr(torch, '__version__', 'unknown')
|
|
torchaudio_version = getattr(torchaudio, '__version__', 'unknown')
|
|
torchvision_version = getattr(torchvision, '__version__', 'unknown')
|
|
|
|
# Detect CUDA vs CPU variants
|
|
torch_cuda = '+cu' in torch_version
|
|
torch_cpu = '+cpu' in torch_version
|
|
torchaudio_cpu = '+cpu' in torchaudio_version or (not '+cu' in torchaudio_version and not torch_cpu)
|
|
torchvision_cpu = '+cpu' in torchvision_version or (not '+cu' in torchvision_version and not torch_cpu)
|
|
|
|
warnings = []
|
|
|
|
# Check for performance-killing mixed installations
|
|
if torch_cuda and torchaudio_cpu:
|
|
warnings.append("❌ CRITICAL PERFORMANCE WARNING: GPU PyTorch + CPU TorchAudio detected!")
|
|
warnings.append(f" PyTorch: {torch_version} | TorchAudio: {torchaudio_version}")
|
|
warnings.append(" This causes major slowdowns and VRAM spikes in TTS generation.")
|
|
warnings.append(" Fix: pip uninstall torchaudio && pip install torchaudio --index-url https://download.pytorch.org/whl/cu129")
|
|
|
|
if torch_cuda and torchvision_cpu:
|
|
warnings.append("❌ PERFORMANCE WARNING: GPU PyTorch + CPU TorchVision detected!")
|
|
warnings.append(f" PyTorch: {torch_version} | TorchVision: {torchvision_version}")
|
|
warnings.append(" Fix: pip uninstall torchvision && pip install torchvision --index-url https://download.pytorch.org/whl/cu129")
|
|
|
|
if warnings:
|
|
print("\n" + "="*80)
|
|
print("🔥 TTS AUDIO SUITE - PYTORCH INSTALLATION WARNING")
|
|
print("="*80)
|
|
for warning in warnings:
|
|
print(warning)
|
|
print("="*80 + "\n")
|
|
|
|
self._pytorch_warning_shown = True
|
|
|
|
except ImportError as e:
|
|
# PyTorch not available - this will be caught elsewhere
|
|
pass
|
|
except Exception as e:
|
|
# Don't let consistency check break model loading
|
|
print(f"⚠️ PyTorch consistency check failed: {e}")
|
|
self._pytorch_warning_shown = True
|
|
|
|
def load_model(self, config: UnifiedModelConfig, force_reload: bool = False) -> Any:
|
|
"""
|
|
Load a model using the unified interface.
|
|
|
|
Args:
|
|
config: Model configuration
|
|
force_reload: Force reload even if cached
|
|
|
|
Returns:
|
|
The loaded model (unwrapped from ComfyUI wrapper)
|
|
|
|
Raises:
|
|
ValueError: If no factory is registered for the engine/model type
|
|
RuntimeError: If model loading fails
|
|
"""
|
|
# Check PyTorch consistency on first model load
|
|
self._check_pytorch_consistency()
|
|
|
|
# Generate unique cache key
|
|
cache_key = self._generate_cache_key(config)
|
|
|
|
# Check if force reload is needed
|
|
if force_reload:
|
|
tts_model_manager.remove_model(cache_key)
|
|
|
|
# Get existing model if cached
|
|
wrapper = tts_model_manager.get_model(cache_key)
|
|
if wrapper is not None:
|
|
# Ensure model is on correct device
|
|
target_device = config.device
|
|
if target_device == "auto":
|
|
target_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# Reload if device mismatch (handles both explicit devices and "auto" resolution)
|
|
if wrapper.current_device != target_device:
|
|
wrapper.model_load(target_device)
|
|
return wrapper.model
|
|
|
|
# Find appropriate factory
|
|
factory_key = f"{config.engine_name}_{config.model_type}"
|
|
if factory_key not in self._model_factories:
|
|
raise ValueError(f"No model factory registered for {factory_key}")
|
|
|
|
factory_func = self._model_factories[factory_key]
|
|
|
|
# Prepare factory arguments (device will be added by model manager)
|
|
factory_kwargs = {
|
|
"model_name": config.model_name,
|
|
"language": config.language,
|
|
"model_path": config.model_path,
|
|
"repo_id": config.repo_id,
|
|
}
|
|
|
|
# Add engine-specific parameters
|
|
if config.additional_params:
|
|
factory_kwargs.update(config.additional_params)
|
|
|
|
# Use ComfyUI model manager to load
|
|
wrapper = tts_model_manager.load_model(
|
|
model_factory_func=factory_func,
|
|
model_key=cache_key,
|
|
model_type=config.model_type,
|
|
engine=config.engine_name,
|
|
device=config.device,
|
|
force_reload=force_reload,
|
|
**factory_kwargs
|
|
)
|
|
|
|
return wrapper.model
|
|
|
|
def unload_model(self, config: UnifiedModelConfig) -> bool:
|
|
"""
|
|
Unload a specific model.
|
|
|
|
Args:
|
|
config: Model configuration to unload
|
|
|
|
Returns:
|
|
True if model was unloaded, False if not found
|
|
"""
|
|
cache_key = self._generate_cache_key(config)
|
|
return tts_model_manager.remove_model(cache_key)
|
|
|
|
def clear_engine_models(self, engine_name: str) -> None:
|
|
"""Clear all models for a specific engine"""
|
|
tts_model_manager.clear_cache(engine=engine_name)
|
|
|
|
def clear_model_type(self, model_type: str) -> None:
|
|
"""Clear all models of a specific type"""
|
|
tts_model_manager.clear_cache(model_type=model_type)
|
|
|
|
def get_model_stats(self) -> Dict[str, Any]:
|
|
"""Get statistics about loaded models"""
|
|
return tts_model_manager.get_stats()
|
|
|
|
def _generate_cache_key(self, config: UnifiedModelConfig) -> str:
|
|
"""Generate unique cache key for model configuration"""
|
|
components = [
|
|
config.engine_name,
|
|
config.model_type,
|
|
config.model_name,
|
|
config.device,
|
|
config.language or "default",
|
|
]
|
|
|
|
# Add path/repo info if available
|
|
if config.model_path:
|
|
components.append(f"path:{Path(config.model_path).name}")
|
|
elif config.repo_id:
|
|
components.append(f"repo:{config.repo_id}")
|
|
|
|
# Add additional_params to ensure attention_mode, quantization, etc. trigger reloads
|
|
if config.additional_params:
|
|
# Sort params for consistent cache keys
|
|
sorted_params = sorted(config.additional_params.items())
|
|
params_str = "_".join([f"{k}:{v}" for k, v in sorted_params])
|
|
components.append(f"params:{params_str}")
|
|
|
|
cache_key = "_".join(components)
|
|
# print(f"🔑 Generated cache key: {cache_key}") # Debug only when needed
|
|
return cache_key
|
|
|
|
|
|
# Global unified interface instance
|
|
unified_model_interface = UnifiedModelInterface()
|
|
|
|
|
|
# Convenience functions for common model operations
|
|
def load_tts_model(engine_name: str,
|
|
model_name: str,
|
|
device: str,
|
|
language: str = "English",
|
|
model_path: Optional[str] = None,
|
|
repo_id: Optional[str] = None,
|
|
force_reload: bool = False,
|
|
**kwargs) -> Any:
|
|
"""
|
|
Convenience function to load TTS models.
|
|
|
|
Args:
|
|
engine_name: Engine name ("chatterbox", "f5tts", etc.)
|
|
model_name: Model identifier
|
|
device: Target device
|
|
language: Model language
|
|
model_path: Local path if available
|
|
repo_id: HuggingFace repo ID if available
|
|
force_reload: Force reload even if cached
|
|
**kwargs: Additional engine-specific parameters
|
|
|
|
Returns:
|
|
Loaded TTS model
|
|
"""
|
|
config = UnifiedModelConfig(
|
|
engine_name=engine_name,
|
|
model_type="tts",
|
|
model_name=model_name,
|
|
device=device,
|
|
language=language,
|
|
model_path=model_path,
|
|
repo_id=repo_id,
|
|
additional_params=kwargs
|
|
)
|
|
return unified_model_interface.load_model(config, force_reload)
|
|
|
|
|
|
def load_vc_model(engine_name: str,
|
|
model_name: str,
|
|
device: str,
|
|
language: str = "English",
|
|
model_path: Optional[str] = None,
|
|
repo_id: Optional[str] = None,
|
|
force_reload: bool = False,
|
|
**kwargs) -> Any:
|
|
"""
|
|
Convenience function to load Voice Conversion models.
|
|
"""
|
|
config = UnifiedModelConfig(
|
|
engine_name=engine_name,
|
|
model_type="vc",
|
|
model_name=model_name,
|
|
device=device,
|
|
language=language,
|
|
model_path=model_path,
|
|
repo_id=repo_id,
|
|
additional_params=kwargs
|
|
)
|
|
return unified_model_interface.load_model(config, force_reload)
|
|
|
|
|
|
def load_auxiliary_model(engine_name: str,
|
|
model_type: str, # "tokenizer", "hubert", "separator", etc.
|
|
model_name: str,
|
|
device: str,
|
|
model_path: Optional[str] = None,
|
|
repo_id: Optional[str] = None,
|
|
force_reload: bool = False,
|
|
**kwargs) -> Any:
|
|
"""
|
|
Convenience function to load auxiliary models (tokenizers, HuBERT, etc.).
|
|
"""
|
|
config = UnifiedModelConfig(
|
|
engine_name=engine_name,
|
|
model_type=model_type,
|
|
model_name=model_name,
|
|
device=device,
|
|
model_path=model_path,
|
|
repo_id=repo_id,
|
|
additional_params=kwargs
|
|
)
|
|
return unified_model_interface.load_model(config, force_reload)
|
|
|
|
|
|
# Factory registration helpers
|
|
def register_chatterbox_factory():
|
|
"""Register ChatterBox model factories with universal component mixing logic"""
|
|
|
|
def load_chatterbox_with_mixing(target_class, device, language, model_path):
|
|
"""Universal ChatterBox loading with component mixing fallback"""
|
|
from engines.chatterbox.language_models import find_local_model_path
|
|
|
|
# Ensure target_class is available before attempting to use it
|
|
if target_class is None:
|
|
raise RuntimeError(f"ChatterBox class not available - cannot load model")
|
|
|
|
# Try provided path first
|
|
if model_path:
|
|
try:
|
|
return target_class.from_local(model_path, device)
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to load from provided path {model_path}: {e}")
|
|
# Continue to fallback logic
|
|
|
|
# Try language-specific local model
|
|
try:
|
|
local_path = find_local_model_path(language)
|
|
if local_path:
|
|
return target_class.from_local(local_path, device)
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to load local {language} model: {e}")
|
|
|
|
# Try HuggingFace download for requested language
|
|
try:
|
|
return target_class.from_pretrained(device, language)
|
|
except Exception as e:
|
|
error_str = str(e)
|
|
print(f"⚠️ Failed to download {language} model: {e}")
|
|
|
|
# Check for 401 authorization errors (gated/private repos)
|
|
if "401" in error_str or "Unauthorized" in error_str:
|
|
print(f"🔒 {language} model requires manual download due to authentication requirements")
|
|
print(f" Please visit the model repository and download manually to use this model")
|
|
|
|
# For German variants, fall back to base German model instead of English
|
|
if language.startswith("German (") and language != "German":
|
|
print(f"🔄 Falling back to base German model for ChatterBox")
|
|
# Try the complete fallback sequence for German
|
|
try:
|
|
german_local = find_local_model_path("German")
|
|
if german_local:
|
|
return target_class.from_local(german_local, device)
|
|
except Exception as local_error:
|
|
print(f"⚠️ German local model failed: {local_error}")
|
|
|
|
# If local German fails, try downloading German
|
|
try:
|
|
print(f"📥 Attempting to download base German model")
|
|
return target_class.from_pretrained(device, "German")
|
|
except Exception as german_error:
|
|
print(f"⚠️ German download also failed: {german_error}")
|
|
# Continue to English fallback
|
|
|
|
# Final fallback: English model if not already trying English
|
|
if language != "English":
|
|
print(f"🔄 Falling back to English model for ChatterBox")
|
|
try:
|
|
# Try English local first
|
|
english_local = find_local_model_path("English")
|
|
if english_local:
|
|
return target_class.from_local(english_local, device)
|
|
else:
|
|
return target_class.from_pretrained(device, "English")
|
|
except Exception as english_error:
|
|
print(f"❌ English fallback also failed: {english_error}")
|
|
raise e # Raise original error
|
|
else:
|
|
raise e # Already trying English, fail
|
|
|
|
def chatterbox_tts_factory(**kwargs):
|
|
try:
|
|
from engines.chatterbox.tts import ChatterboxTTS
|
|
except ImportError:
|
|
ChatterboxTTS = None
|
|
|
|
if ChatterboxTTS is None:
|
|
raise RuntimeError("ChatterboxTTS not available - check installation")
|
|
|
|
device = kwargs.get("device", "auto")
|
|
language = kwargs.get("language", "English")
|
|
model_path = kwargs.get("model_path")
|
|
|
|
return load_chatterbox_with_mixing(ChatterboxTTS, device, language, model_path)
|
|
|
|
def chatterbox_vc_factory(**kwargs):
|
|
try:
|
|
from engines.chatterbox.vc import ChatterboxVC
|
|
from engines.chatterbox.language_models import supports_voice_conversion, get_vc_supported_languages
|
|
except ImportError:
|
|
ChatterboxVC = None
|
|
supports_voice_conversion = None
|
|
|
|
if ChatterboxVC is None:
|
|
raise RuntimeError("ChatterboxVC not available - check installation or add bundled version")
|
|
|
|
device = kwargs.get("device", "auto")
|
|
language = kwargs.get("language", "English")
|
|
model_path = kwargs.get("model_path")
|
|
|
|
# Check if language supports VC before attempting to load
|
|
if supports_voice_conversion and not supports_voice_conversion(language):
|
|
print(f"❌ {language} model does not support voice conversion")
|
|
print(f" Voice conversion requires s3gen component which is missing from this model")
|
|
print(f" Try English model (confirmed working) or German models (tested working)")
|
|
print(f" Other language models may not have VC components")
|
|
|
|
raise RuntimeError(f"Voice conversion not supported for {language} model. "
|
|
f"Try English or German models instead.")
|
|
|
|
return load_chatterbox_with_mixing(ChatterboxVC, device, language, model_path)
|
|
|
|
unified_model_interface.register_model_factory("chatterbox", "tts", chatterbox_tts_factory)
|
|
unified_model_interface.register_model_factory("chatterbox", "vc", chatterbox_vc_factory)
|
|
|
|
|
|
def register_f5tts_factory():
|
|
"""Register F5-TTS model factories"""
|
|
def f5tts_factory(**kwargs):
|
|
from engines.f5tts.f5tts import ChatterBoxF5TTS
|
|
from utils.models.smart_loader import smart_model_loader
|
|
|
|
device = kwargs.get("device", "auto")
|
|
model_name = kwargs.get("model_name", "F5TTS_Base")
|
|
|
|
# Use Smart Loader for consistency with the base node approach
|
|
def f5tts_load_callback(device: str, model: str) -> Any:
|
|
"""Factory callback for F5-TTS model loading"""
|
|
# Try local first, then HuggingFace
|
|
try:
|
|
# Check for local models
|
|
import os
|
|
import folder_paths
|
|
search_paths = [
|
|
os.path.join(folder_paths.models_dir, "TTS", "F5-TTS", model),
|
|
os.path.join(folder_paths.models_dir, "F5-TTS", model), # Legacy
|
|
os.path.join(folder_paths.models_dir, "Checkpoints", "F5-TTS", model) # Legacy
|
|
]
|
|
|
|
local_path = None
|
|
for path in search_paths:
|
|
if os.path.exists(path):
|
|
local_path = path
|
|
break
|
|
|
|
if local_path:
|
|
print(f"📁 Using local F5-TTS model: {local_path}")
|
|
return ChatterBoxF5TTS.from_local(local_path, device, model)
|
|
else:
|
|
print(f"📥 Loading F5-TTS model '{model}' from HuggingFace")
|
|
return ChatterBoxF5TTS.from_pretrained(device, model)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load F5-TTS model '{model}': {e}")
|
|
|
|
model, _ = smart_model_loader.load_model_if_needed(
|
|
engine_type="f5tts",
|
|
model_name=model_name,
|
|
current_model=None, # Factory always loads fresh
|
|
device=device,
|
|
load_callback=f5tts_load_callback,
|
|
force_reload=False
|
|
)
|
|
|
|
return model
|
|
|
|
unified_model_interface.register_model_factory("f5tts", "tts", f5tts_factory)
|
|
|
|
|
|
def register_higgs_audio_factory():
|
|
"""Register Higgs Audio model factories with stateless wrapper for ComfyUI integration"""
|
|
def higgs_audio_factory(**kwargs):
|
|
from engines.higgs_audio.boson_multimodal.serve.serve_engine import HiggsAudioServeEngine
|
|
from engines.higgs_audio.higgs_audio import HiggsAudioEngine
|
|
from engines.higgs_audio.stateless_wrapper import create_stateless_higgs_wrapper
|
|
|
|
model_name = kwargs.get("model_name")
|
|
device = kwargs.get("device", "cuda")
|
|
model_path = kwargs.get("model_path", model_name)
|
|
tokenizer_path = kwargs.get("tokenizer_path")
|
|
enable_cuda_graphs = kwargs.get("enable_cuda_graphs", True)
|
|
|
|
# Check if we need to force recreate due to CUDA Graph corruption
|
|
force_recreate = kwargs.get("_force_recreate", False)
|
|
if force_recreate:
|
|
print(f"🔥 Force recreating Higgs Audio engine due to CUDA Graph corruption")
|
|
|
|
# Log CUDA Graph setting
|
|
perf_mode = "High Performance (CUDA Graphs ON)" if enable_cuda_graphs else "Memory Safe (CUDA Graphs OFF)"
|
|
print(f"🔧 Higgs Audio mode: {perf_mode}")
|
|
|
|
# Create the base HiggsAudioEngine but initialize manually to avoid circular dependency
|
|
base_engine = HiggsAudioEngine()
|
|
|
|
# Manually set up the engine properties to avoid calling initialize_engine
|
|
base_engine.model_path = model_path
|
|
base_engine.tokenizer_path = tokenizer_path
|
|
base_engine.device = device
|
|
|
|
# Get smart paths using the engine's methods
|
|
model_path_for_engine = base_engine._get_smart_model_path(model_path)
|
|
tokenizer_path_for_engine = base_engine._get_smart_tokenizer_path(tokenizer_path)
|
|
|
|
# Configure CUDA Graph settings based on toggle
|
|
if enable_cuda_graphs:
|
|
print(f"🚀 Loading Higgs Audio model directly on {device} for optimal performance...")
|
|
cache_type = "StaticCache"
|
|
kv_cache_lengths = [1024, 2048, 4096] # StaticCache for CUDA Graph optimization
|
|
else:
|
|
print(f"🚀 Loading Higgs Audio model directly on {device} (CUDA Graphs disabled for memory safety)...")
|
|
cache_type = "DynamicCache"
|
|
# Still need cache buckets but they won't use StaticCache
|
|
kv_cache_lengths = [1024, 2048, 4096] # Same sizes but will use DynamicCache
|
|
|
|
# Disable CUDA Graph compilation at environment level
|
|
import os
|
|
os.environ['PYTORCH_DISABLE_CUDA_GRAPH'] = '1'
|
|
os.environ['TORCH_COMPILE_DISABLE'] = '1'
|
|
|
|
serve_engine = HiggsAudioServeEngine(
|
|
model_name_or_path=model_path_for_engine,
|
|
audio_tokenizer_name_or_path=tokenizer_path_for_engine,
|
|
device=device,
|
|
kv_cache_lengths=kv_cache_lengths,
|
|
enable_cuda_graphs=enable_cuda_graphs
|
|
)
|
|
|
|
# Settings are already stored by the constructor
|
|
|
|
print(f"✅ Higgs Audio model loaded directly on {device}")
|
|
|
|
# Set the serve engine on the base engine
|
|
base_engine.engine = serve_engine
|
|
|
|
# Wrap with stateless wrapper for ComfyUI model management compatibility
|
|
stateless_wrapper = create_stateless_higgs_wrapper(base_engine)
|
|
|
|
print(f"📦 Higgs Audio engine ready (via unified interface)")
|
|
return stateless_wrapper
|
|
|
|
unified_model_interface.register_model_factory("higgs_audio", "tts", higgs_audio_factory)
|
|
|
|
|
|
def register_rvc_factory():
|
|
"""Register RVC model factories"""
|
|
def rvc_factory(**kwargs):
|
|
"""Load RVC model and wrap it for ComfyUI management"""
|
|
from engines.rvc.impl.vc_infer_pipeline import get_vc
|
|
from engines.rvc.impl.config import config as rvc_config
|
|
from engines.rvc.rvc_engine import RVCModelWrapper
|
|
|
|
model_path = kwargs.get("model_path")
|
|
index_path = kwargs.get("index_path")
|
|
device = kwargs.get("device", "cuda")
|
|
|
|
if not model_path:
|
|
raise ValueError("RVC factory requires model_path")
|
|
|
|
# Load the RVC model using the actual implementation
|
|
model_data = get_vc(model_path, index_path, rvc_config, device)
|
|
|
|
# Wrap the model_data dict so it has .to() method for device management
|
|
wrapped_model = RVCModelWrapper(model_data, device)
|
|
|
|
return wrapped_model
|
|
|
|
def hubert_factory(**kwargs):
|
|
from engines.rvc.impl.lib.model_utils import load_hubert
|
|
model_path = kwargs.get("model_path")
|
|
config = kwargs.get("config", {})
|
|
return load_hubert(model_path, config)
|
|
|
|
def mdx23c_separator_factory(**kwargs):
|
|
from engines.rvc.impl.lib.mdx23c import MDX23C
|
|
model_path = kwargs.get("model_path")
|
|
device = kwargs.get("device", "cuda")
|
|
return MDX23C(model_path, device)
|
|
|
|
def demucs_separator_factory(**kwargs):
|
|
from engines.rvc.impl.lib.uvr5_pack.demucs.pretrained import get_model
|
|
model_name = kwargs.get("model_name", "htdemucs")
|
|
device = kwargs.get("device", "cuda")
|
|
return get_model(model_name).to(device)
|
|
|
|
unified_model_interface.register_model_factory("rvc", "vc", rvc_factory)
|
|
unified_model_interface.register_model_factory("rvc", "hubert", hubert_factory)
|
|
unified_model_interface.register_model_factory("rvc", "separator_mdx", mdx23c_separator_factory)
|
|
unified_model_interface.register_model_factory("rvc", "separator_demucs", demucs_separator_factory)
|
|
|
|
|
|
def register_vibevoice_factory():
|
|
"""Register VibeVoice model factory"""
|
|
def vibevoice_factory(**kwargs):
|
|
"""Factory for VibeVoice models with ComfyUI integration"""
|
|
|
|
# Ensure accelerate is imported before VibeVoice engine
|
|
try:
|
|
import accelerate
|
|
# (accelerate version check removed - only show on errors)
|
|
except ImportError as e:
|
|
print(f"⚠️ Unified Interface: accelerate not available: {e}")
|
|
|
|
from engines.vibevoice_engine.vibevoice_engine import VibeVoiceEngine
|
|
|
|
# Extract parameters
|
|
model_name = kwargs.get("model_name", "vibevoice-1.5B")
|
|
device = kwargs.get("device", "auto")
|
|
attention_mode = kwargs.get("attention_mode", "auto")
|
|
quantize_llm_4bit = kwargs.get("quantize_llm_4bit", False)
|
|
|
|
# Create engine instance
|
|
engine = VibeVoiceEngine()
|
|
|
|
# Initialize with parameters
|
|
engine.initialize_engine(
|
|
model_name=model_name,
|
|
device=device,
|
|
attention_mode=attention_mode,
|
|
quantize_llm_4bit=quantize_llm_4bit
|
|
)
|
|
|
|
print(f"✅ VibeVoice model '{model_name}' loaded via unified interface")
|
|
|
|
# Return the engine which contains both model and processor
|
|
return engine
|
|
|
|
unified_model_interface.register_model_factory("vibevoice", "tts", vibevoice_factory)
|
|
|
|
|
|
def register_index_tts_factory():
|
|
"""Register IndexTTS-2 model factory"""
|
|
def index_tts_factory(**kwargs):
|
|
"""Factory for IndexTTS-2 models with ComfyUI integration"""
|
|
import os
|
|
import sys
|
|
|
|
# Extract parameters
|
|
model_path = kwargs.get("model_path")
|
|
device = kwargs.get("device", "auto")
|
|
use_fp16 = kwargs.get("use_fp16", True)
|
|
use_cuda_kernel = kwargs.get("use_cuda_kernel", None)
|
|
use_deepspeed = kwargs.get("use_deepspeed", False)
|
|
|
|
if not model_path or not os.path.exists(model_path):
|
|
raise RuntimeError(f"IndexTTS-2 model not found at {model_path}. Auto-download should have been triggered earlier.")
|
|
|
|
try:
|
|
# Add bundled IndexTTS path to sys.path so internal imports work
|
|
import sys
|
|
bundled_path = os.path.join(os.path.dirname(__file__), "..", "..", "engines", "index_tts")
|
|
bundled_path = os.path.abspath(bundled_path)
|
|
if bundled_path not in sys.path:
|
|
sys.path.insert(0, bundled_path)
|
|
|
|
# Check for conflicting external IndexTTS installation
|
|
try:
|
|
import indextts
|
|
external_path = indextts.__file__
|
|
if ('site-packages' in external_path or
|
|
'conda' in external_path or
|
|
'pip' in external_path or
|
|
'dist-packages' in external_path):
|
|
|
|
raise ImportError(f"""
|
|
❌ External IndexTTS installation detected!
|
|
|
|
Found at: {external_path}
|
|
|
|
This conflicts with our bundled IndexTTS-2 engine and causes import errors.
|
|
|
|
🔧 SOLUTION: Please uninstall the external version:
|
|
pip uninstall indextts
|
|
|
|
Then restart ComfyUI.
|
|
|
|
Our bundled version has all required dependencies and will work perfectly.
|
|
""")
|
|
except ImportError as e:
|
|
if "External IndexTTS" in str(e):
|
|
raise e
|
|
# No external installation found - this is good!
|
|
pass
|
|
except Exception:
|
|
# Some other issue with indextts, proceed anyway
|
|
pass
|
|
|
|
# Import from our bundled IndexTTS engine
|
|
from engines.index_tts.indextts.infer_v2 import IndexTTS2
|
|
|
|
# Initialize IndexTTS-2 engine
|
|
config_path = os.path.join(model_path, "config.yaml")
|
|
|
|
# Verify config file exists after download
|
|
if not os.path.exists(config_path):
|
|
raise RuntimeError(f"IndexTTS-2 config.yaml not found at {config_path} even after download. Please check model integrity.")
|
|
|
|
engine = IndexTTS2(
|
|
cfg_path=config_path,
|
|
model_dir=model_path,
|
|
device=device,
|
|
use_fp16=use_fp16 and device != "cpu",
|
|
use_cuda_kernel=use_cuda_kernel,
|
|
use_deepspeed=use_deepspeed
|
|
)
|
|
|
|
print(f"✅ IndexTTS-2 model loaded via unified interface on {device}")
|
|
return engine
|
|
|
|
except ImportError as e:
|
|
raise ImportError(f"IndexTTS-2 dependencies not available. Error: {e}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load IndexTTS-2 model: {e}")
|
|
|
|
unified_model_interface.register_model_factory("index_tts", "tts", index_tts_factory)
|
|
|
|
|
|
def register_chatterbox_23lang_factory():
|
|
"""Register ChatterBox Official 23-Lang model factory"""
|
|
def chatterbox_23lang_factory(**kwargs):
|
|
"""Factory for ChatterBox Official 23-Lang models with ComfyUI integration"""
|
|
from engines.chatterbox_official_23lang.tts import ChatterboxOfficial23LangTTS
|
|
import folder_paths
|
|
import os
|
|
import torch
|
|
|
|
# Extract parameters
|
|
device = kwargs.get("device", "auto")
|
|
model_name = kwargs.get("model_name", "Official 23-Lang") # Always same model
|
|
language = kwargs.get("language", "english") # This is the actual language to use
|
|
model_version = kwargs.get("model_version", "v2") # v1 or v2
|
|
|
|
# Determine device
|
|
if device == "auto":
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# Get model directory path
|
|
models_dir = folder_paths.models_dir
|
|
ckpt_dir = os.path.join(models_dir, "TTS", "chatterbox_official_23lang", "Official 23-Lang")
|
|
|
|
print(f"🌍 Loading ChatterBox Official 23-Lang model for {language} on {device}")
|
|
print(f"📁 Using model directory: {ckpt_dir}")
|
|
|
|
# Try local first, then use from_pretrained for auto-download if needed
|
|
if os.path.exists(ckpt_dir):
|
|
try:
|
|
engine = ChatterboxOfficial23LangTTS.from_local(
|
|
ckpt_dir=ckpt_dir,
|
|
device=device,
|
|
model_name="Official 23-Lang",
|
|
model_version=model_version
|
|
)
|
|
print(f"✅ ChatterBox Official 23-Lang '{language}' loaded via unified interface")
|
|
return engine
|
|
except FileNotFoundError as e:
|
|
print(f"⚠️ Local model incomplete: {e}")
|
|
print(f"📥 Downloading missing {model_version} files...")
|
|
|
|
# Use from_pretrained for auto-download
|
|
engine = ChatterboxOfficial23LangTTS.from_pretrained(
|
|
device=device,
|
|
model_name="ChatterBox Official 23-Lang",
|
|
model_version=model_version
|
|
)
|
|
|
|
print(f"✅ ChatterBox Official 23-Lang '{language}' loaded via unified interface")
|
|
|
|
# Return the engine
|
|
return engine
|
|
|
|
unified_model_interface.register_model_factory("chatterbox_official_23lang", "tts", chatterbox_23lang_factory)
|
|
|
|
|
|
def initialize_all_factories():
|
|
"""Initialize all model factories"""
|
|
register_chatterbox_factory()
|
|
register_chatterbox_23lang_factory()
|
|
register_f5tts_factory()
|
|
register_higgs_audio_factory()
|
|
register_rvc_factory()
|
|
register_vibevoice_factory()
|
|
register_index_tts_factory()
|
|
|
|
|
|
# Auto-initialize on import
|
|
initialize_all_factories() |