iGavroche
c989a48aab
- LTX Video VAE detection (128 channels, 32x spatial compression, 8x temporal) - z-image / z-image-turbo pixel-space VAE passthrough (no decode needed) - fp16 default on all AMD GPUs (Strix Halo, RDNA 2/3, CDNA) - Per-architecture tuning with automatic detection (gfx1151/gfx1100/gfx1030/CDNA) - Temporal tiling for long videos: overlapping chunks with linear blend, zero frame loss, progress bar in UI (disabled by default) - Chunked IO bypass for LTX — fixes OOM with max_split_size_mb caps - Fixed memory estimation: correct for all dims (including 5D temporal) and dtypes - APU-aware memory management (psutil for unified memory on Strix Halo) - Tile size floor for high-compression VAEs (min 64 latent pixels) - 25 unit tests covering all new features
142 lines
4.1 KiB
Python
142 lines
4.1 KiB
Python
"""
|
|
Memory management utilities for ROCM Ninodes.
|
|
|
|
Provides functions for GPU memory cleanup, monitoring, and safety checks
|
|
optimized for ROCm and AMD GPUs, including APU (unified memory) awareness.
|
|
"""
|
|
|
|
import gc
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
|
|
import torch
|
|
|
|
|
|
def simple_memory_cleanup() -> bool:
|
|
"""Simple memory cleanup for ROCm"""
|
|
try:
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
print(f" Memory cleanup error: {e}")
|
|
return False
|
|
|
|
|
|
def gentle_memory_cleanup() -> bool:
|
|
"""Gentle memory cleanup - less aggressive for mature ROCm drivers"""
|
|
try:
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
gc.collect()
|
|
return True
|
|
except Exception as e:
|
|
print(f" Memory cleanup error: {e}")
|
|
return False
|
|
|
|
|
|
def aggressive_memory_cleanup() -> bool:
|
|
"""Gentle memory cleanup - renamed for compatibility"""
|
|
return gentle_memory_cleanup()
|
|
|
|
|
|
def emergency_memory_cleanup() -> bool:
|
|
"""Emergency memory cleanup - gentle approach for mature ROCm"""
|
|
try:
|
|
if not torch.cuda.is_available():
|
|
return False
|
|
|
|
torch.cuda.synchronize()
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
return True
|
|
except Exception as e:
|
|
logging.error(f"Emergency memory cleanup failed: {e}")
|
|
return False
|
|
|
|
|
|
def get_gpu_memory_info(is_apu: bool = False) -> Tuple[Optional[int], Optional[int], Optional[int], Optional[int]]:
|
|
"""Get accurate GPU memory information
|
|
|
|
For discrete GPUs: returns VRAM stats from torch.cuda.
|
|
For APUs (unified memory): returns system-wide available memory
|
|
since GPU and CPU share the same pool.
|
|
|
|
Args:
|
|
is_apu: Whether the GPU is an APU with unified memory (e.g. Strix Halo)
|
|
|
|
Returns:
|
|
Tuple of (total_memory, allocated_memory, reserved_memory, free_memory)
|
|
Returns (None, None, None, None) if CUDA is not available
|
|
"""
|
|
if not torch.cuda.is_available():
|
|
return None, None, None, None
|
|
|
|
if is_apu:
|
|
total_memory = torch.cuda.get_device_properties(0).total_memory
|
|
allocated_memory = torch.cuda.memory_allocated(0)
|
|
reserved_memory = torch.cuda.memory_reserved(0)
|
|
import psutil
|
|
system = psutil.virtual_memory()
|
|
free_memory = int(system.available)
|
|
return total_memory, allocated_memory, reserved_memory, free_memory
|
|
|
|
total_memory = torch.cuda.get_device_properties(0).total_memory
|
|
allocated_memory = torch.cuda.memory_allocated(0)
|
|
reserved_memory = torch.cuda.memory_reserved(0)
|
|
free_memory = total_memory - reserved_memory
|
|
|
|
return total_memory, allocated_memory, reserved_memory, free_memory
|
|
|
|
|
|
def check_memory_safety(required_memory_gb: float = 2.0, is_apu: bool = False) -> Tuple[bool, str]:
|
|
"""Check if there's enough memory for the operation
|
|
|
|
For APUs, checks against system-wide available memory.
|
|
For discrete GPUs, checks against VRAM.
|
|
|
|
Args:
|
|
required_memory_gb: Required memory in GB
|
|
is_apu: Whether the GPU is an APU with unified memory
|
|
|
|
Returns:
|
|
Tuple of (is_safe, message)
|
|
"""
|
|
if not torch.cuda.is_available():
|
|
return True, "CUDA not available"
|
|
|
|
_, _, _, free_memory = get_gpu_memory_info(is_apu=is_apu)
|
|
|
|
if free_memory is None:
|
|
return True, "Memory info unavailable"
|
|
|
|
free_gb = free_memory / (1024**3)
|
|
required_gb = required_memory_gb
|
|
|
|
if free_gb < required_gb:
|
|
return False, f"Only {free_gb:.2f}GB free, need {required_gb:.2f}GB"
|
|
|
|
return True, f"Memory OK: {free_gb:.2f}GB free"
|
|
|
|
|
|
def is_apu_architecture(arch_name: str) -> bool:
|
|
"""Detect if the GPU architecture is an APU (unified memory)"""
|
|
apu_archs = ["gfx1151", "gfx1150"]
|
|
return any(a in arch_name for a in apu_archs)
|
|
|
|
|
|
__all__ = [
|
|
'simple_memory_cleanup',
|
|
'gentle_memory_cleanup',
|
|
'aggressive_memory_cleanup',
|
|
'emergency_memory_cleanup',
|
|
'get_gpu_memory_info',
|
|
'check_memory_safety',
|
|
'is_apu_architecture',
|
|
]
|