269 lines
7.6 KiB
Python
269 lines
7.6 KiB
Python
"""
|
|
Device capabilities and memory management utilities.
|
|
|
|
Integrates with ComfyUI's model_management for safe memory operations.
|
|
"""
|
|
import os
|
|
import gc
|
|
import psutil
|
|
import torch
|
|
from typing import Optional, Tuple
|
|
from dataclasses import dataclass
|
|
|
|
# Import ComfyUI's model management functions
|
|
try:
|
|
from comfy.model_management import (
|
|
get_torch_device,
|
|
get_free_memory,
|
|
get_total_memory,
|
|
free_memory,
|
|
soft_empty_cache,
|
|
unload_all_models,
|
|
is_device_cuda,
|
|
)
|
|
COMFY_AVAILABLE = True
|
|
except ImportError:
|
|
COMFY_AVAILABLE = False
|
|
|
|
|
|
@dataclass
|
|
class DeviceCapabilities:
|
|
"""Device capabilities for I/O and compute optimization."""
|
|
device_type: str # 'cuda', 'cpu', 'mps', 'xpu'
|
|
total_vram_gb: float
|
|
free_vram_gb: float
|
|
total_ram_gb: float
|
|
free_ram_gb: float
|
|
cpu_cores: int
|
|
recommended_workers: int
|
|
max_tensor_size_gb: float # Largest tensor we can safely load
|
|
supports_pinned_memory: bool
|
|
supports_async_streams: bool
|
|
|
|
|
|
def get_device_capabilities(device: Optional[torch.device] = None) -> DeviceCapabilities:
|
|
"""
|
|
Analyze device capabilities for optimal I/O configuration.
|
|
|
|
Args:
|
|
device: Target device (default: ComfyUI's torch device)
|
|
|
|
Returns:
|
|
DeviceCapabilities dataclass with recommended settings
|
|
"""
|
|
if device is None:
|
|
if COMFY_AVAILABLE:
|
|
device = get_torch_device()
|
|
else:
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
device_type = device.type if hasattr(device, 'type') else 'cpu'
|
|
|
|
# Get RAM info
|
|
ram_info = psutil.virtual_memory()
|
|
total_ram_gb = ram_info.total / (1024**3)
|
|
free_ram_gb = ram_info.available / (1024**3)
|
|
|
|
# Get VRAM info
|
|
if device_type == 'cuda' and torch.cuda.is_available():
|
|
if COMFY_AVAILABLE:
|
|
total_vram = get_total_memory(device)
|
|
free_vram = get_free_memory(device)
|
|
else:
|
|
free_vram, total_vram = torch.cuda.mem_get_info(device)
|
|
total_vram_gb = total_vram / (1024**3)
|
|
free_vram_gb = free_vram / (1024**3)
|
|
supports_pinned = True
|
|
supports_async = True
|
|
elif device_type == 'mps':
|
|
total_vram_gb = total_ram_gb * 0.5 # MPS shares RAM
|
|
free_vram_gb = free_ram_gb * 0.5
|
|
supports_pinned = False
|
|
supports_async = False
|
|
else:
|
|
total_vram_gb = 0
|
|
free_vram_gb = 0
|
|
supports_pinned = False
|
|
supports_async = False
|
|
|
|
# CPU cores
|
|
cpu_cores = os.cpu_count() or 4
|
|
|
|
# Recommended workers based on hardware
|
|
if device_type == 'cuda':
|
|
# GPU: fewer workers since I/O is not the bottleneck
|
|
recommended_workers = min(4, cpu_cores // 2)
|
|
else:
|
|
# CPU: more workers for parallel I/O
|
|
recommended_workers = min(8, cpu_cores)
|
|
|
|
# Max tensor size: leave headroom for operations
|
|
if device_type == 'cuda':
|
|
max_tensor_gb = min(free_vram_gb * 0.8, free_ram_gb * 0.5)
|
|
else:
|
|
max_tensor_gb = free_ram_gb * 0.6
|
|
|
|
return DeviceCapabilities(
|
|
device_type=device_type,
|
|
total_vram_gb=total_vram_gb,
|
|
free_vram_gb=free_vram_gb,
|
|
total_ram_gb=total_ram_gb,
|
|
free_ram_gb=free_ram_gb,
|
|
cpu_cores=cpu_cores,
|
|
recommended_workers=recommended_workers,
|
|
max_tensor_size_gb=max_tensor_gb,
|
|
supports_pinned_memory=supports_pinned,
|
|
supports_async_streams=supports_async,
|
|
)
|
|
|
|
|
|
def estimate_model_size(file_path: str) -> float:
|
|
"""Estimate model size in GB from file size."""
|
|
try:
|
|
return os.path.getsize(file_path) / (1024**3)
|
|
except OSError:
|
|
return 0.0
|
|
|
|
|
|
def can_fit_in_memory(
|
|
model_paths: list[str],
|
|
device: Optional[torch.device] = None,
|
|
safety_factor: float = 0.8
|
|
) -> Tuple[bool, str]:
|
|
"""
|
|
Check if models can fit in available memory.
|
|
|
|
Args:
|
|
model_paths: List of model file paths
|
|
device: Target device
|
|
safety_factor: Fraction of free memory to use (default 0.8)
|
|
|
|
Returns:
|
|
(can_fit, reason_message)
|
|
"""
|
|
caps = get_device_capabilities(device)
|
|
|
|
total_size = sum(estimate_model_size(p) for p in model_paths)
|
|
|
|
if caps.device_type == 'cuda':
|
|
available = caps.free_vram_gb * safety_factor
|
|
mem_type = "VRAM"
|
|
else:
|
|
available = caps.free_ram_gb * safety_factor
|
|
mem_type = "RAM"
|
|
|
|
if total_size <= available:
|
|
return True, f"OK: {total_size:.2f}GB needed, {available:.2f}GB {mem_type} available"
|
|
else:
|
|
return False, f"Insufficient memory: {total_size:.2f}GB needed but only {available:.2f}GB {mem_type} available"
|
|
|
|
|
|
def prepare_for_large_operation(
|
|
estimated_memory_gb: float,
|
|
device: Optional[torch.device] = None
|
|
) -> bool:
|
|
"""
|
|
Prepare device for a large memory operation by freeing resources.
|
|
|
|
Args:
|
|
estimated_memory_gb: Estimated memory needed in GB
|
|
device: Target device
|
|
|
|
Returns:
|
|
True if preparation succeeded
|
|
"""
|
|
if not COMFY_AVAILABLE:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
return True
|
|
|
|
if device is None:
|
|
device = get_torch_device()
|
|
|
|
memory_bytes = int(estimated_memory_gb * 1024**3)
|
|
|
|
# Use ComfyUI's free_memory which unloads models intelligently
|
|
try:
|
|
free_memory(memory_bytes, device)
|
|
return True
|
|
except Exception as e:
|
|
print(f"[DeviceUtils] Warning: Could not free memory: {e}")
|
|
return False
|
|
|
|
|
|
def cleanup_after_operation():
|
|
"""Clean up memory after a large operation."""
|
|
gc.collect()
|
|
|
|
if COMFY_AVAILABLE:
|
|
soft_empty_cache()
|
|
elif torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.ipc_collect()
|
|
|
|
|
|
def get_optimal_workers(
|
|
model_size_gb: float,
|
|
caps: Optional[DeviceCapabilities] = None
|
|
) -> int:
|
|
"""
|
|
Calculate optimal number of I/O workers based on model size and device.
|
|
|
|
Args:
|
|
model_size_gb: Model size in GB
|
|
caps: Pre-calculated device capabilities (optional)
|
|
|
|
Returns:
|
|
Recommended number of workers
|
|
"""
|
|
if caps is None:
|
|
caps = get_device_capabilities()
|
|
|
|
# For very large models, fewer workers to avoid memory pressure
|
|
if model_size_gb > caps.free_ram_gb * 0.5:
|
|
return max(1, caps.recommended_workers // 2)
|
|
|
|
# For small models, use more workers
|
|
if model_size_gb < 2.0:
|
|
return min(8, caps.cpu_cores)
|
|
|
|
return caps.recommended_workers
|
|
|
|
|
|
def get_processing_device_string(prefer_cuda: bool = True) -> str:
|
|
"""
|
|
Get the best available processing device as a string.
|
|
|
|
Args:
|
|
prefer_cuda: Prefer CUDA over CPU if available
|
|
|
|
Returns:
|
|
Device string ('cuda', 'mps', 'cpu')
|
|
"""
|
|
if prefer_cuda:
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
|
return "mps"
|
|
return "cpu"
|
|
|
|
|
|
def print_device_info():
|
|
"""Print device capabilities for debugging."""
|
|
caps = get_device_capabilities()
|
|
print(f"\n{'='*50}")
|
|
print(" Device Capabilities")
|
|
print(f"{'='*50}")
|
|
print(f" Device Type: {caps.device_type}")
|
|
print(f" Total VRAM: {caps.total_vram_gb:.2f} GB")
|
|
print(f" Free VRAM: {caps.free_vram_gb:.2f} GB")
|
|
print(f" Total RAM: {caps.total_ram_gb:.2f} GB")
|
|
print(f" Free RAM: {caps.free_ram_gb:.2f} GB")
|
|
print(f" CPU Cores: {caps.cpu_cores}")
|
|
print(f" Recommended Workers: {caps.recommended_workers}")
|
|
print(f" Max Tensor Size: {caps.max_tensor_size_gb:.2f} GB")
|
|
print(f" Pinned Memory: {caps.supports_pinned_memory}")
|
|
print(f" Async Streams: {caps.supports_async_streams}")
|
|
print(f"{'='*50}\n")
|