277 lines
11 KiB
Python
277 lines
11 KiB
Python
"""
|
|
F5-TTS Model Manager Extension
|
|
Extends the existing ModelManager for F5-TTS specific functionality
|
|
"""
|
|
|
|
import os
|
|
import warnings
|
|
import torch
|
|
import folder_paths
|
|
from typing import Optional, List, Tuple, Dict, Any
|
|
|
|
# Import unified model interface for ComfyUI integration
|
|
from utils.models.unified_model_interface import load_tts_model
|
|
|
|
|
|
class F5TTSModelManager:
|
|
"""
|
|
F5-TTS model manager extending the existing ModelManager pattern.
|
|
Handles F5-TTS specific model discovery and loading.
|
|
"""
|
|
|
|
# Class-level cache for F5-TTS models
|
|
_f5tts_model_cache: Dict[str, Any] = {}
|
|
_f5tts_model_sources: Dict[str, str] = {}
|
|
|
|
def __init__(self, node_dir: Optional[str] = None):
|
|
"""
|
|
Initialize F5TTSModelManager.
|
|
|
|
Args:
|
|
node_dir: Optional override for the node directory path
|
|
"""
|
|
self.node_dir = node_dir or os.path.dirname(os.path.dirname(__file__))
|
|
self.f5tts_available = False
|
|
self._check_f5tts_availability()
|
|
|
|
def _check_f5tts_availability(self):
|
|
"""Check if F5-TTS is available"""
|
|
try:
|
|
import engines.f5_tts
|
|
self.f5tts_available = True
|
|
except ImportError:
|
|
self.f5tts_available = False
|
|
|
|
def find_f5tts_models(self, model_name: str = None) -> List[Tuple[str, Optional[str]]]:
|
|
"""
|
|
Find F5-TTS model files in order of priority:
|
|
1. ComfyUI models/F5-TTS/ directory (primary location)
|
|
2. ComfyUI models/Checkpoints/F5-TTS/ directory (fallback for user convenience)
|
|
3. HuggingFace download
|
|
|
|
Args:
|
|
model_name: Optional specific model name to search for
|
|
|
|
Returns:
|
|
List of tuples containing (source_type, path) in priority order
|
|
"""
|
|
model_paths = []
|
|
|
|
# Search paths in order of priority: TTS first, then legacy
|
|
search_paths = [
|
|
os.path.join(folder_paths.models_dir, "TTS", "F5-TTS"),
|
|
os.path.join(folder_paths.models_dir, "F5-TTS"), # Legacy
|
|
os.path.join(folder_paths.models_dir, "Checkpoints", "F5-TTS") # Legacy
|
|
]
|
|
|
|
# 1-2. Check ComfyUI models folders
|
|
for search_path in search_paths:
|
|
if os.path.exists(search_path):
|
|
for item in os.listdir(search_path):
|
|
item_path = os.path.join(search_path, item)
|
|
if os.path.isdir(item_path):
|
|
# If specific model name provided, only check matching folders
|
|
if model_name and item.lower() != model_name.lower():
|
|
continue
|
|
|
|
# Check if it contains model files
|
|
has_model = False
|
|
for ext in [".safetensors", ".pt"]:
|
|
model_files = [f for f in os.listdir(item_path) if f.endswith(ext)]
|
|
if model_files:
|
|
has_model = True
|
|
break
|
|
if has_model:
|
|
model_paths.append(("comfyui", item_path))
|
|
|
|
# 3. HuggingFace download as fallback
|
|
model_paths.append(("huggingface", None))
|
|
|
|
return model_paths
|
|
|
|
def get_f5tts_model_cache_key(self, model_name: str, device: str, source: str, path: Optional[str] = None) -> str:
|
|
"""
|
|
Generate a cache key for F5-TTS model instances.
|
|
|
|
Args:
|
|
model_name: Name of the F5-TTS model
|
|
device: Target device ('cuda', 'cpu')
|
|
source: Model source ('comfyui', 'huggingface')
|
|
path: Optional path for local models
|
|
|
|
Returns:
|
|
Cache key string
|
|
"""
|
|
path_component = path or "default"
|
|
return f"f5tts_{model_name}_{device}_{source}_{path_component}"
|
|
|
|
def load_f5tts_model(self, model_name: str = "F5TTS_Base", device: str = "auto",
|
|
force_reload: bool = False) -> Any:
|
|
"""
|
|
Load F5-TTS model with ComfyUI-integrated caching support.
|
|
|
|
Args:
|
|
model_name: Name of the F5-TTS model to load
|
|
device: Target device ('auto', 'cuda', 'cpu')
|
|
force_reload: Force reload even if cached
|
|
|
|
Returns:
|
|
F5-TTS model instance
|
|
|
|
Raises:
|
|
ImportError: If F5-TTS is not available
|
|
RuntimeError: If model loading fails
|
|
"""
|
|
if not self.f5tts_available:
|
|
raise ImportError("F5-TTS not available - check installation")
|
|
|
|
# Resolve auto device
|
|
if device == "auto":
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# Try unified model interface for ComfyUI integration
|
|
try:
|
|
model = load_tts_model(
|
|
engine_name="f5tts",
|
|
model_name=model_name,
|
|
device=device,
|
|
force_reload=force_reload
|
|
)
|
|
|
|
return model
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to load F5-TTS model via unified interface: {e}")
|
|
print(f"🔄 Falling back to direct loading...")
|
|
# Fall back to original logic if unified interface fails
|
|
|
|
# Get available model paths for specific model
|
|
model_paths = self.find_f5tts_models(model_name)
|
|
|
|
model_loaded = False
|
|
last_error = None
|
|
|
|
# Check if this model should be available locally
|
|
local_model_expected = False
|
|
f5tts_search_paths = [
|
|
os.path.join(folder_paths.models_dir, "TTS", "F5-TTS"),
|
|
os.path.join(folder_paths.models_dir, "F5-TTS"), # Legacy
|
|
os.path.join(folder_paths.models_dir, "Checkpoints", "F5-TTS") # Legacy
|
|
]
|
|
|
|
for f5tts_path in f5tts_search_paths:
|
|
if os.path.exists(f5tts_path):
|
|
for item in os.listdir(f5tts_path):
|
|
if item.lower() == model_name.lower() or model_name.lower() in item.lower():
|
|
local_model_expected = True
|
|
break
|
|
if local_model_expected:
|
|
break
|
|
|
|
for source, path in model_paths:
|
|
try:
|
|
cache_key = self.get_f5tts_model_cache_key(model_name, device, source, path)
|
|
|
|
# Check class-level cache first
|
|
if not force_reload and cache_key in self._f5tts_model_cache:
|
|
model = self._f5tts_model_cache[cache_key]
|
|
self._f5tts_model_sources[cache_key] = source
|
|
model_loaded = True
|
|
break
|
|
|
|
# Load model based on source
|
|
if source == "comfyui" and path:
|
|
from engines.f5tts.f5tts import ChatterBoxF5TTS
|
|
model = ChatterBoxF5TTS.from_local(path, device, model_name)
|
|
elif source == "huggingface":
|
|
# If a local model was expected but failed, don't fallback to HuggingFace
|
|
if local_model_expected and last_error:
|
|
print(f"❌ Model '{model_name}' found in local directory but failed to load")
|
|
print(f"❌ Fallback to HuggingFace disabled for locally expected models")
|
|
raise RuntimeError(f"Model '{model_name}' exists locally but failed to load. Please check the model files or use a different model.")
|
|
|
|
# Check if model is available on HuggingFace
|
|
# This prevents calling from_pretrained with unsupported models,
|
|
# which would trigger engine fallback and potentially bypass cache
|
|
from engines.f5tts.f5tts import F5TTS_MODELS
|
|
if model_name not in F5TTS_MODELS:
|
|
print(f"❌ Model '{model_name}' not available on HuggingFace")
|
|
raise RuntimeError(f"Model '{model_name}' is not available on HuggingFace. Please install it locally or use a different model.")
|
|
|
|
from engines.f5tts.f5tts import ChatterBoxF5TTS
|
|
model = ChatterBoxF5TTS.from_pretrained(device, model_name)
|
|
else:
|
|
continue
|
|
|
|
# Cache the loaded model
|
|
self._f5tts_model_cache[cache_key] = model
|
|
self._f5tts_model_sources[cache_key] = source
|
|
model_loaded = True
|
|
break
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ Failed to load F5-TTS model from {source}: {str(e)}")
|
|
last_error = e
|
|
continue
|
|
|
|
if not model_loaded:
|
|
error_msg = f"Failed to load F5-TTS model '{model_name}' from any source"
|
|
if last_error:
|
|
error_msg += f". Last error: {last_error}"
|
|
raise RuntimeError(error_msg)
|
|
|
|
return self._f5tts_model_cache[cache_key]
|
|
|
|
def get_f5tts_model_source(self, model_name: str, device: str) -> Optional[str]:
|
|
"""
|
|
Get the source of a cached F5-TTS model.
|
|
|
|
Args:
|
|
model_name: Name of the F5-TTS model
|
|
device: Device the model is loaded on
|
|
|
|
Returns:
|
|
Model source string or None if not cached
|
|
"""
|
|
# Try to find in cache by checking different sources with their paths
|
|
model_paths = self.find_f5tts_models(model_name)
|
|
|
|
for source, path in model_paths:
|
|
cache_key = self.get_f5tts_model_cache_key(model_name, device, source, path)
|
|
if cache_key in self._f5tts_model_sources:
|
|
return self._f5tts_model_sources[cache_key]
|
|
|
|
return None
|
|
|
|
def clear_f5tts_cache(self):
|
|
"""Clear F5-TTS model cache."""
|
|
self._f5tts_model_cache.clear()
|
|
self._f5tts_model_sources.clear()
|
|
|
|
def get_f5tts_model_configs(self) -> List[str]:
|
|
"""
|
|
Get available F5-TTS model configurations.
|
|
|
|
Returns:
|
|
List of available F5-TTS model names
|
|
"""
|
|
try:
|
|
from engines.f5tts.f5tts import get_f5tts_models
|
|
return get_f5tts_models()
|
|
except ImportError:
|
|
# Fallback list if F5-TTS not available
|
|
return ["F5TTS_Base", "F5TTS_v1_Base", "E2TTS_Base"]
|
|
|
|
@property
|
|
def is_f5tts_available(self) -> bool:
|
|
"""
|
|
Check if F5-TTS is available.
|
|
|
|
Returns:
|
|
True if F5-TTS is available, False otherwise
|
|
"""
|
|
return self.f5tts_available
|
|
|
|
|
|
# Global F5-TTS model manager instance
|
|
f5tts_model_manager = F5TTSModelManager() |