diodiogod
9f2166a297
Fix critical device handling issues that cause cache mismatches and unnecessary reloads: 1. Cache 'auto' device resolution (torch_device_resolver.py) - 'auto' device now cached after first resolution - Prevents cache key instability when device availability changes mid-session - Add reset_device_cache() for manual reset if needed - Fixes: GPU memory pressure causing unexpected CPU fallback 2. Fix device comparison in unified interface (unified_model_interface.py) - Resolve both target and wrapper device before comparison - Handle both string and torch.device object representations - Fixes: Device type mismatch causing models to stay on wrong device 3. Consolidate device resolution (base_node.py) - Replace duplicate resolve_device() with centralized torch_device_resolver - Now supports MPS (Apple Silicon), CUDA, XPU, and CPU consistently - Fixes: Missing MPS/XPU detection in base node Impact: - Eliminates unnecessary model reloads in long sessions - Consistent device handling across all node types - Better GPU memory management with stable cache keys
92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
"""
|
|
Torch Device Resolver - Intelligent device selection for TTS engines
|
|
Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU with automatic detection
|
|
"""
|
|
|
|
import torch
|
|
from typing import Literal
|
|
|
|
# Cache for "auto" device resolution to ensure consistency across calls
|
|
# Critical for cache key generation in model loading
|
|
_cached_auto_device = None
|
|
|
|
|
|
def resolve_torch_device(device: str) -> str:
|
|
"""
|
|
Resolve device string to actual available device on current system.
|
|
|
|
Handles intelligent device selection with fallback chain:
|
|
- "auto": Detects best available (MPS > CUDA > XPU > CPU) and caches result
|
|
- "mps": Apple Silicon GPU (macOS only)
|
|
- "cuda": NVIDIA GPU
|
|
- "xpu": Intel GPU
|
|
- "cpu": CPU fallback
|
|
|
|
Args:
|
|
device: Device specification string ("auto", "cuda", "mps", "xpu", or "cpu")
|
|
|
|
Returns:
|
|
Resolved device string: "cuda", "mps", "xpu", or "cpu"
|
|
|
|
Example:
|
|
>>> device = resolve_torch_device("auto") # Returns "cuda" on Linux with NVIDIA
|
|
>>> device = resolve_torch_device("auto") # Returns "mps" on macOS with Apple Silicon
|
|
>>> device = resolve_torch_device("auto") # Returns "xpu" on Linux with Intel GPU
|
|
>>> device = resolve_torch_device("auto") # Returns "cpu" on any system without GPU
|
|
|
|
Note:
|
|
The "auto" device resolution is cached after first call to ensure consistency.
|
|
This prevents cache key mismatches in model loading when device availability changes mid-session.
|
|
"""
|
|
global _cached_auto_device
|
|
|
|
if device == "auto":
|
|
# Return cached resolution if available to ensure consistency
|
|
if _cached_auto_device is not None:
|
|
return _cached_auto_device
|
|
|
|
# Intelligent device detection with proper fallback chain
|
|
# Check MPS first (Apple Silicon)
|
|
if torch.backends.mps.is_available():
|
|
_cached_auto_device = "mps"
|
|
# Check CUDA second (NVIDIA GPUs)
|
|
elif torch.cuda.is_available():
|
|
_cached_auto_device = "cuda"
|
|
# Check XPU third (Intel GPUs)
|
|
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
|
_cached_auto_device = "xpu"
|
|
# CPU fallback (always available)
|
|
else:
|
|
_cached_auto_device = "cpu"
|
|
|
|
return _cached_auto_device
|
|
|
|
# Explicit device selection - pass through as-is
|
|
# Validation happens when torch.device() is created
|
|
return device
|
|
|
|
|
|
def reset_device_cache() -> None:
|
|
"""
|
|
Reset cached device resolution.
|
|
|
|
Use this if device availability changes (e.g., after GPU reset or power state change).
|
|
Normally should not be needed - cache persists for session consistency.
|
|
"""
|
|
global _cached_auto_device
|
|
_cached_auto_device = None
|
|
|
|
|
|
def validate_device_string(device: str) -> bool:
|
|
"""
|
|
Validate that device string is valid.
|
|
|
|
Args:
|
|
device: Device string to validate
|
|
|
|
Returns:
|
|
True if device is valid, False otherwise
|
|
"""
|
|
valid_devices = {"auto", "cuda", "mps", "xpu", "cpu"}
|
|
return device.lower() in valid_devices
|