Files

T

diodiogod 43d1342b4b Fix: Move TIER2-3_REFACTOR_PLAN.md to correct directory

Remove incorrectly escaped 'Dev\ reports' directory

2025-11-14 23:02:11 -03:00

6.8 KiB

Raw Permalink Blame History

Model Loading Tier 2-3 Refactor Plan

Overview

Complete the model loading architecture refactor by:

Tier 2: Consolidate caching, standardize error handling
Tier 3: Standardize factory signatures, implement capability registry

TIER 2: Consolidation & Standardization

2.1 Remove SmartModelLoader Usage

Current State: SmartModelLoader still used in:

nodes/chatterbox/chatterbox_tts_node.py (lines with smart_model_loader)
engines/chatterbox/streaming_model_manager.py
engines/chatterbox_official_23lang/streaming_model_manager.py
engines/adapters/f5tts_streaming_adapter.py

Why Remove:

Creates third caching system (after unified interface + ComfyUI manager)
Causes cache synchronization issues
Delegates to manager.load_tts_model() anyway (redundant)

Approach:

ChatterBox TTS node: Replace smart_model_loader call with direct unified interface
Streaming processors: Use unified interface for model loading
Verify no other code depends on SmartModelLoader returns

Risk: Medium - streaming processors have complex state management

2.2 Standardize Error Handling

Current State: 5 different error patterns across factories

Target State: Single exception hierarchy:

class TTSModelLoadingError(Exception):
    """Base exception for model loading"""
    pass

class TTSModelNotFoundError(TTSModelLoadingError):
    """Model not found locally or on HuggingFace"""
    pass

class TTSDeviceError(TTSModelLoadingError):
    """Device-related issues"""
    pass

class TTSModelFormatError(TTSModelLoadingError):
    """Model format compatibility issues"""
    pass

Files to Update:

Create utils/models/exceptions.py with exception classes
Update all factories to use consistent exceptions
Update node error handlers to catch standardized exceptions

Risk: Low - backwards compatible

2.3 Unified Cache Key Generation

Current State: Multiple cache key formats:

Unified interface: string concatenation
SmartModelLoader: MD5 hashing
Nodes: MD5 of dict items

Target State: Single hash-based cache key function

def generate_model_cache_key(
    engine_name: str,
    model_name: str,
    device: str,
    language: Optional[str] = None,
    **extra_params
) -> str:
    """Generate standardized cache key"""

Files to Create/Update:

utils/models/cache_key_generator.py (new)
Update unified_model_interface.py to use it
Update all node caching to use it

Risk: Low

TIER 3: Architecture Standardization

3.1 Factory Parameter Dataclass

Current State: Each factory has different parameters

Target State: All factories accept standardized dataclass:

@dataclass
class ModelLoadConfig:
    """Standard configuration for all model factories"""
    # Required
    device: str  # "auto", "cuda", "cpu", etc.
    model_name: str

    # Optional (engine-specific)
    language: Optional[str] = None
    model_path: Optional[str] = None
    repo_id: Optional[str] = None

    # Additional params (for quantization, optimization, etc.)
    additional_params: Optional[Dict[str, Any]] = None

    def __post_init__(self):
        # Resolve device before storing
        from utils.device import resolve_torch_device
        self.device = resolve_torch_device(self.device)

Factory Signature Becomes:

def engine_factory(config: ModelLoadConfig) -> Any:
    """All factories use same signature"""

Files to Create/Update:

utils/models/factory_config.py (new)
All factories in unified_model_interface.py
Unified interface load_model() method

Risk: High - changes all factory signatures

3.2 Engine Capability Registry

Current State: Hardcoded engine-specific logic in model_manager.py

Target State: Data-driven engine capabilities:

ENGINE_CAPABILITIES = {
    "chatterbox": {
        "supports_voice_conversion": True,
        "has_fallback": True,
        "can_corrupt": False,
        "recovery_handler": None
    },
    "higgs_audio": {
        "supports_voice_conversion": False,
        "has_fallback": False,
        "can_corrupt": True,
        "recovery_handler": reset_higgs_cuda_graphs,
        "requires_special_init": True
    },
    # ... all 7 engines
}

Files to Create:

utils/models/engine_registry.py (new) - define capabilities
Remove hardcoded engine checks from model_manager.py
Use registry for generic recovery

Risk: Medium - affects error recovery paths

3.3 Unified Fallback Handler

Current State: ChatterBox has sophisticated fallback, others don't

Target State: Generic fallback utility any engine can use:

class ModelFallbackHandler:
    """Generic fallback chain for model loading"""

    def __init__(self, primary_loader: Callable):
        self.primary_loader = primary_loader
        self.fallback_chain = []

    def add_fallback(self, fallback_loader: Callable, condition: Callable):
        """Add fallback with condition"""
        pass

    def load(self) -> Any:
        """Execute fallback chain"""
        pass