119 lines
4.8 KiB
Python
119 lines
4.8 KiB
Python
import torch
|
|
import gc
|
|
import comfy.model_management as mm
|
|
from . import register_node
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
any_type = AnyType("*")
|
|
|
|
def format_bytes(bytes_val):
|
|
return f"{bytes_val / (1024 ** 3):.2f} GB"
|
|
|
|
@register_node
|
|
class MeisoftVRAMDefragmenter:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"anything": (any_type,),
|
|
"force_model_unload": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (any_type, "STRING")
|
|
RETURN_NAMES = ("anything", "log")
|
|
FUNCTION = "defragment"
|
|
CATEGORY = "🔁 Sequential Batcher/Tools"
|
|
|
|
def defragment(self, anything, force_model_unload):
|
|
log_output = []
|
|
def _log(msg):
|
|
print(msg)
|
|
log_output.append(str(msg))
|
|
_log(f"\n{'='*65}")
|
|
_log(f"🧹 [Secuencial Batcher] NODO: VRAM Defragmenter (Modo Arquitectura Profunda)")
|
|
|
|
# --- ESTADÍSTICAS ANTES ---
|
|
alloc_before = 0
|
|
res_before = 0
|
|
if torch.cuda.is_available():
|
|
alloc_before = torch.cuda.memory_allocated()
|
|
res_before = torch.cuda.memory_reserved()
|
|
_log(f" -> 📊 Estado Inicial | Asignada: {format_bytes(alloc_before)} | Reservada (Caché): {format_bytes(res_before)}")
|
|
|
|
# --- FASE 1: CONFIGURACIÓN DINÁMICA DEL ALLOCATOR ---
|
|
try:
|
|
# Permite que la reserva de VRAM crezca en bloque elástico
|
|
torch.cuda.memory._set_allocator_settings("expandable_segments:True")
|
|
_log(" -> ⚙️ PyTorch Allocator ajustado a expandable_segments:True")
|
|
except Exception as e:
|
|
_log(f" -> ⚠️ Aviso: No se pudo ajustar expandable_segments: {e}")
|
|
|
|
# --- FASE 2: EVACUACIÓN (Opcional) ---
|
|
if force_model_unload:
|
|
_log(" -> 📦 Trasladando modelos inactivos a la RAM (unload_all_models)...")
|
|
mm.unload_all_models()
|
|
|
|
# --- FASE 3: LIMPIEZA INTERNA DE COMFYUI (LA CLAVE) ---
|
|
_log(" -> 🧹 Destruyendo referencias ocultas en el gestor de ComfyUI...")
|
|
if hasattr(mm, "cleanup_models"):
|
|
mm.cleanup_models()
|
|
if hasattr(mm, "cleanup_models_gc"):
|
|
mm.cleanup_models_gc()
|
|
|
|
# --- FASE 4: ROMPER REFERENCIAS PYTHON ---
|
|
_log(" -> 🗑️ Forzando Garbage Collector de Python (Doble pasada)...")
|
|
gc.collect()
|
|
gc.collect() # Doble pasada para romper referencias circulares rebeldes
|
|
|
|
# --- FASE 5: EVACUACIÓN SUAVE ---
|
|
mm.soft_empty_cache()
|
|
|
|
# --- FASE 6: PURGADO PROFUNDO DE PYTORCH ---
|
|
if torch.cuda.is_available():
|
|
_log(" -> 🧱 Desfragmentando Caching Allocator de CUDA...")
|
|
torch.cuda.ipc_collect()
|
|
torch.cuda.synchronize() # 🚀 OPTIMIZACIÓN: Esperar a que terminen los kernels pendientes PRIMERO
|
|
torch.cuda.empty_cache() # 🚀 Luego vaciamos la caché de forma segura
|
|
|
|
# --- FASE 7: LIBERACIÓN DE MEMORIA DEL SISTEMA (Multi-Plataforma) ---
|
|
import platform
|
|
try:
|
|
import ctypes
|
|
if platform.system() == "Windows":
|
|
# API de Windows: Reduce el 'Working Set' del proceso actual al mínimo
|
|
kernel32 = ctypes.windll.kernel32
|
|
current_process = kernel32.GetCurrentProcess()
|
|
kernel32.SetProcessWorkingSetSize(current_process, -1, -1)
|
|
_log(" -> 💻 Working Set reducido (RAM del sistema purgada en Windows).")
|
|
else:
|
|
# API de Linux/WSL: Libera la memoria de la librería C
|
|
ctypes.CDLL('libc.so.6').malloc_trim(0)
|
|
_log(" -> 💻 Malloc Trim ejecutado (RAM del sistema purgada en Linux/WSL).")
|
|
except Exception as e:
|
|
# Si algo falla a nivel de SO, fallamos silenciosamente sin romper el nodo
|
|
pass
|
|
|
|
# --- ESTADÍSTICAS DESPUÉS ---
|
|
if torch.cuda.is_available():
|
|
alloc_after = torch.cuda.memory_allocated()
|
|
res_after = torch.cuda.memory_reserved()
|
|
|
|
# Cálculos de liberación
|
|
freed_alloc = alloc_before - alloc_after
|
|
freed_res = res_before - res_after
|
|
|
|
sign_alloc = "+" if freed_alloc < 0 else "-"
|
|
sign_res = "+" if freed_res < 0 else "-"
|
|
|
|
_log(f" -> 📈 Estado Final | Asignada: {format_bytes(alloc_after)} | Reservada (Caché): {format_bytes(res_after)}")
|
|
_log(f" -> ✨ TOTAL LIBERADO | {sign_alloc}{abs(freed_alloc)/(1024**3):.2f} GB (Asignada) | {sign_res}{abs(freed_res)/(1024**3):.2f} GB (Reservada purgada)")
|
|
else:
|
|
_log(" -> ✨ Limpieza completada (Modo CPU/Otra arquitectura).")
|
|
|
|
_log(f"{'='*65}\n")
|
|
return (anything, "\n".join(log_output))
|