From f06ebf70d46c0bd8b804926eba5401b0c9355dde Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 10 Jul 2025 09:36:35 -0400 Subject: [PATCH 1/4] Disabled tensor counting in BlockSwap's debug mode Disabled tensor counting in BlockSwap's debug mode (show_tensors=False) to prevent conflicts with other ComfyUI extensions. The tensor counting feature was triggering lazy initialization of unrelated libraries (like OpenAI) causing 'api_key' errors. --- src/interfaces/comfyui_node.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/interfaces/comfyui_node.py b/src/interfaces/comfyui_node.py index 092b036..b1ee431 100644 --- a/src/interfaces/comfyui_node.py +++ b/src/interfaces/comfyui_node.py @@ -230,7 +230,7 @@ class SeedVR2: # BlockSwap debugger memory state if debugger is not None: cleanup_stage = "partial" if should_keep_model else "full" - debugger.log_memory_state(f"After {cleanup_stage} cleanup", show_tensors=True) + debugger.log_memory_state(f"After {cleanup_stage} cleanup", show_tensors=False) def _internal_execute(self, images, model, seed, new_resolution, cfg_scale, batch_size, preserve_vram, temporal_overlap, debug, block_swap_config): From 8924245de20ac5cc0e7d96c6f41dad3e986e3535 Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 10 Jul 2025 09:36:36 -0400 Subject: [PATCH 2/4] Disabled tensor counting in BlockSwap's debug mode Disabled tensor counting in BlockSwap's debug mode (show_tensors=False) to prevent conflicts with other ComfyUI extensions. The tensor counting feature was triggering lazy initialization of unrelated libraries (like OpenAI) causing 'api_key' errors. --- src/optimization/blockswap.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/optimization/blockswap.py b/src/optimization/blockswap.py index e7c6f02..546a731 100644 --- a/src/optimization/blockswap.py +++ b/src/optimization/blockswap.py @@ -193,7 +193,7 @@ def apply_block_swap_to_dit(runner, block_swap_config: Dict[str, Any]) -> None: debugger.log(f"Configuring: {blocks_to_swap}/{total_blocks} blocks for swapping") - debugger.log_memory_state("Before BlockSwap", show_tensors=True) + debugger.log_memory_state("Before BlockSwap", show_tensors=False) # Configure I/O components offload_io_components = block_swap_config.get("offload_io_components", False) @@ -235,7 +235,7 @@ def apply_block_swap_to_dit(runner, block_swap_config: Dict[str, Any]) -> None: # Protect model from being moved entirely _protect_model_from_move(model, runner, debugger) - debugger.log_memory_state("After BlockSwap", show_tensors=True) + debugger.log_memory_state("After BlockSwap", show_tensors=False) debugger.log("✅ BlockSwap configuration complete") From 0d55c70d3d13090d8d9fa370df75e2246d809aac Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 10 Jul 2025 09:36:40 -0400 Subject: [PATCH 3/4] Disabled tensor counting in BlockSwap's debug mode Disabled tensor counting in BlockSwap's debug mode (show_tensors=False) to prevent conflicts with other ComfyUI extensions. The tensor counting feature was triggering lazy initialization of unrelated libraries (like OpenAI) causing 'api_key' errors. --- src/core/generation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/core/generation.py b/src/core/generation.py index cddc857..3f4baf2 100644 --- a/src/core/generation.py +++ b/src/core/generation.py @@ -441,7 +441,7 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si #clear_vram_cache() # Log memory state at the end of each batch if debugger: - debugger.log_memory_state(f"Batch {batch_number} - Memory", show_tensors=True) + debugger.log_memory_state(f"Batch {batch_number} - Memory", show_tensors=False) finally: if debugger: @@ -458,7 +458,7 @@ def generation_loop(runner, images, cfg_scale=1.0, seed=666, res_w=720, batch_si # Log final memory state if debugger: - debugger.log_memory_state("Generation loop - After cleanup", show_tensors=True) + debugger.log_memory_state("Generation loop - After cleanup", show_tensors=False) # OPTIMISATION ULTIME : Pré-allocation et copie directe (évite les torch.cat multiples) print(f"💾 Processing {len(batch_samples)} batch_samples with memory-optimized pre-allocation") From 68f7e913dfe1b4e872bd896765e6e3660ef7d083 Mon Sep 17 00:00:00 2001 From: Adrien Toupet Date: Thu, 10 Jul 2025 13:43:34 -0400 Subject: [PATCH 4/4] Avoid debug logging/operations with debug false --- src/optimization/blockswap.py | 60 +++++++++++++++--------------- src/optimization/memory_manager.py | 2 - 2 files changed, 30 insertions(+), 32 deletions(-) diff --git a/src/optimization/blockswap.py b/src/optimization/blockswap.py index 546a731..480231d 100644 --- a/src/optimization/blockswap.py +++ b/src/optimization/blockswap.py @@ -76,15 +76,14 @@ class BlockSwapDebugger: component_type: Type of component ("block" or "io") direction: Direction of swap ("compute" or "offload") """ - # Store timing data with component info - self.swap_times.append({ - 'component_id': component_id, - 'component_type': component_type, - 'duration': duration, - 'direction': direction - }) - if self.enabled: + # Store timing data with component info + self.swap_times.append({ + 'component_id': component_id, + 'component_type': component_type, + 'duration': duration, + 'direction': direction + }) # Format message based on component type if component_type == "block": message = f"Block {component_id} swap {direction}: {duration*1000:.1f}ms" @@ -97,31 +96,32 @@ class BlockSwapDebugger: def log_memory_state(self, stage: str, show_tensors: bool = False) -> None: """Log current memory state for debugging.""" - # GPU Memory - if torch.cuda.is_available(): - allocated_gb, reserved_gb, peak_gb = get_vram_usage() - vram_info = f"VRAM: {allocated_gb:.2f}/{reserved_gb:.2f}GB (peak: {peak_gb:.2f}GB)" - self.vram_history.append(allocated_gb) - else: - vram_info = "VRAM: CPU mode" + if self.enabled: + # GPU Memory + if torch.cuda.is_available(): + allocated_gb, reserved_gb, peak_gb = get_vram_usage() + vram_info = f"VRAM: {allocated_gb:.2f}/{reserved_gb:.2f}GB (peak: {peak_gb:.2f}GB)" + self.vram_history.append(allocated_gb) + else: + vram_info = "VRAM: CPU mode" - # RAM Memory - ram_info = "" - if psutil: - try: - process = psutil.Process() - ram_gb = process.memory_info().rss / (1024**3) - ram_info = f" | RAM: {ram_gb:.1f}GB" - except Exception: - pass + # RAM Memory + ram_info = "" + if psutil: + try: + process = psutil.Process() + ram_gb = process.memory_info().rss / (1024**3) + ram_info = f" | RAM: {ram_gb:.1f}GB" + except Exception: + pass - # Tensor count (optional - expensive operation) - tensor_info = "" - if show_tensors: - tensor_count = sum(1 for obj in gc.get_objects() if torch.is_tensor(obj)) - tensor_info = f" | Tensors: {tensor_count}" + # Tensor count (optional - expensive operation) + tensor_info = "" + if show_tensors: + tensor_count = sum(1 for obj in gc.get_objects() if torch.is_tensor(obj)) + tensor_info = f" | Tensors: {tensor_count}" - self.log(f"🧮 {stage}: {vram_info}{ram_info}{tensor_info}") + self.log(f"🧮 {stage}: {vram_info}{ram_info}{tensor_info}") def clear_history(self) -> None: """Clear accumulated history.""" diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index 3c59565..c9031a6 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -271,8 +271,6 @@ def clear_all_caches(runner, debugger=None) -> int: def log_message(message, level="INFO"): if debugger and debugger.enabled: debugger.log(message, level) - else: - print(f" {message}") cleaned_items = 0