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
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
"""
|
|
Base class for engine-specific handlers
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Optional, TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from ..base_wrapper import ComfyUIModelWrapper
|
|
|
|
|
|
class BaseEngineHandler(ABC):
|
|
"""
|
|
Base class for engine-specific model handling.
|
|
|
|
Provides interface for engine-specific unloading and memory management.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def partially_unload(self, wrapper: 'ComfyUIModelWrapper', device: str, memory_to_free: int) -> int:
|
|
"""
|
|
Partially unload the model to free memory.
|
|
|
|
Args:
|
|
wrapper: The ComfyUI model wrapper
|
|
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
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def model_unload(self, wrapper: 'ComfyUIModelWrapper', memory_to_free: Optional[int], unpatch_weights: bool) -> bool:
|
|
"""
|
|
Fully unload the model from GPU memory.
|
|
|
|
Args:
|
|
wrapper: The ComfyUI model wrapper
|
|
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
|
|
"""
|
|
pass |