Merge pull request #38 from AInVFX/blockswap

Preventing Blockswap from lazy-loading unrelated libraries (like OpenAI) causing 'api_key' errors.

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")

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):

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."""
@@ -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")
 
 

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