From 6d8c278a4e21cb1ba42e83ebf72039a1dea262d9 Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Tue, 29 Jul 2025 23:59:20 +0200 Subject: [PATCH] Fix QC bugs post updates --- src/interfaces/comfyui_node.py | 5 --- src/optimization/blockswap.py | 1 - src/optimization/memory_manager.py | 60 ++++++------------------------ src/utils/constants.py | 2 +- src/utils/model_registry.py | 2 +- 5 files changed, 13 insertions(+), 57 deletions(-) diff --git a/src/interfaces/comfyui_node.py b/src/interfaces/comfyui_node.py index 206da35..b989ac5 100644 --- a/src/interfaces/comfyui_node.py +++ b/src/interfaces/comfyui_node.py @@ -303,11 +303,6 @@ class SeedVR2: if force_ram_cleanup: fast_ram_cleanup() - # Clear any remaining references in the class - for attr in list(self.__dict__.keys()): - if attr not in ['__class__', '__module__', '__qualname__']: - setattr(self, attr, None) - def _internal_execute(self, images, model, seed, new_resolution, cfg_scale, batch_size, preserve_vram, temporal_overlap, cache_model, block_swap_config): diff --git a/src/optimization/blockswap.py b/src/optimization/blockswap.py index 29ee2f3..21c1327 100644 --- a/src/optimization/blockswap.py +++ b/src/optimization/blockswap.py @@ -17,7 +17,6 @@ import types import torch import weakref import gc - from typing import Dict, Any, List, Tuple, Optional, Union from src.optimization.memory_manager import get_vram_usage from src.optimization.compatibility import call_rope_with_stability diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index 75bb8ea..dc82142 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -180,62 +180,32 @@ def clear_rope_lru_caches(model) -> int: def fast_model_cleanup(model): - """Aggressive model cleanup that actually frees memory""" + """Fast model cleanup without logs""" if model is None: return - # First move to CPU if on GPU - if next(model.parameters(), None) is not None: - device = next(model.parameters()).device - if device.type == 'cuda': - model.to("cpu") + # Move to CPU + model.to("cpu") - # Aggressively clear all parameters and buffers + # Clear parameters and buffers recursively def clear_recursive(m): - # Clear children first for child in m.children(): clear_recursive(child) - - # Delete all parameters - for name, param in list(m.named_parameters(recurse=False)): + for param in m.parameters(): if param is not None: - # Release the underlying storage - if param.data.numel() > 0: - param.data.set_() # This releases the storage + param.data = param.data.cpu() param.grad = None - delattr(m, name.split('.')[-1]) - - # Delete all buffers - for name, buffer in list(m.named_buffers(recurse=False)): + for buffer in m.buffers(): if buffer is not None: - # Release the underlying storage - if buffer.numel() > 0: - buffer.set_() # This releases the storage - delattr(m, name.split('.')[-1]) - - # Clear any other tensor attributes - for attr_name in list(vars(m).keys()): - attr = getattr(m, attr_name, None) - if torch.is_tensor(attr): - if attr.numel() > 0: - attr.set_() # Release storage - delattr(m, attr_name) + buffer.data = buffer.data.cpu() clear_recursive(model) - - # Clear the module dict to break circular references - model._modules.clear() def fast_ram_cleanup(): - """Aggressive RAM cleanup""" - # Clear Python's internal caches - import sys - sys.intern.clear() if hasattr(sys.intern, 'clear') else None - - # Multiple aggressive garbage collection passes - for _ in range(3): - gc.collect(2) # Collect all generations + """Fast RAM cleanup without excessive logging""" + # Garbage collection + gc.collect() # Clear CUDA cache if torch.cuda.is_available(): @@ -249,14 +219,6 @@ def fast_ram_cleanup(): except: pass - # Force Python to release memory back to OS (Linux) - try: - import ctypes - libc = ctypes.CDLL("libc.so.6") - libc.malloc_trim(0) - except: - pass - def clear_all_caches(runner, debug) -> int: """ diff --git a/src/utils/constants.py b/src/utils/constants.py index e13dd2f..30704ac 100644 --- a/src/utils/constants.py +++ b/src/utils/constants.py @@ -4,6 +4,7 @@ Only includes constants actually used in the codebase """ import os +import folder_paths # Model folder names SEEDVR2_FOLDER_NAME = "SEEDVR2" # Physical folder name on disk @@ -20,7 +21,6 @@ def get_script_directory(): def get_base_cache_dir(): """Get or create the model cache directory""" try: - import folder_paths cache_dir = os.path.join(folder_paths.models_dir, SEEDVR2_FOLDER_NAME) folder_paths.add_model_folder_path(SEEDVR2_MODEL_TYPE, cache_dir) except: diff --git a/src/utils/model_registry.py b/src/utils/model_registry.py index 6b1410c..fe2f365 100644 --- a/src/utils/model_registry.py +++ b/src/utils/model_registry.py @@ -6,6 +6,7 @@ Central registry for model definitions, repositories, and metadata from typing import Dict, List, Optional from dataclasses import dataclass from src.utils.constants import SEEDVR2_MODEL_TYPE, is_supported_model_file, get_base_cache_dir +import folder_paths @dataclass class ModelInfo: @@ -54,7 +55,6 @@ def get_available_models() -> List[str]: model_list = get_default_models() try: - import folder_paths # Ensure the folder is registered before trying to list files get_base_cache_dir() # Get all models from the SEEDVR2 folder using centralized constant