From fc29f8753f960666ab57b2e50b2734b80add3be6 Mon Sep 17 00:00:00 2001
From: Jason Kende <jk@samizdat.ink>
Date: Fri, 13 Jun 2025 01:11:39 -0400
Subject: [PATCH] Fix: Clear model cache in wrappers to prevent CUDA device
 mismatch after long sessions

Adds model cache clearing and preprocessor instance reset before execution to prevent stale cached models with wrong device assignments from causing "Expected all tensors to be on the same device" errors in MultiGPU setups after extended runtime.
---
 nodes.py | 54 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 10 deletions(-)

diff --git a/nodes.py b/nodes.py
index e7e516c..49b7114 100644
--- a/nodes.py
+++ b/nodes.py
@@ -8,19 +8,20 @@ logger = logging.getLogger(__name__)
 class MultiGPUPreprocessorWrapper:
     """
     Base wrapper class that temporarily overrides device placement during preprocessor model loading
-    to prevent multi-GPU device conflicts in ControlNet preprocessors.
+    AND execution to prevent multi-GPU device conflicts in ControlNet preprocessors.
     
     The problem: Preprocessors auto-download models and load them using model_management.get_torch_device(),
     but ComfyUI-MultiGPU monkey-patches this function with dynamic device assignment, causing model 
     components to split across devices and trigger "Expected all tensors to be on the same device" errors.
     
-    The solution: Temporarily override get_torch_device() to return consistent device (cuda:0) during
-    model loading, then restore normal MultiGPU behavior.
+    The solution: Override get_torch_device() to return consistent device (cuda:0) during
+    the ENTIRE preprocessor execution (loading + inference), then restore normal MultiGPU behavior.
     """
     
     def __init__(self, preprocessor_class, target_device='cuda:0'):
         self.preprocessor_class = preprocessor_class
         self.target_device = target_device
+        self._preprocessor_instance = None
         
     @classmethod
     def INPUT_TYPES(cls):
@@ -36,25 +37,58 @@ class MultiGPUPreprocessorWrapper:
     FUNCTION = "execute"
     CATEGORY = "preprocessors/gpu_wrapper"
     
+    def _ensure_tensors_on_device(self, obj, device):
+        """
+        Recursively move tensors in nested data structures to the specified device.
+        """
+        if isinstance(obj, torch.Tensor):
+            return obj.to(device)
+        elif isinstance(obj, dict):
+            return {k: self._ensure_tensors_on_device(v, device) for k, v in obj.items()}
+        elif isinstance(obj, (list, tuple)):
+            return type(obj)(self._ensure_tensors_on_device(item, device) for item in obj)
+        else:
+            return obj
+    
     def execute(self, **kwargs):
         """
-        Execute preprocessor with temporary device override to prevent multi-GPU conflicts.
+        Execute preprocessor with extended device override to prevent multi-GPU conflicts.
+        Now covers both model loading AND inference execution.
         """
         # Critical: Save original function
         original_get_device = model_management.get_torch_device
         
         try:
-            # Temporarily override with consistent device
-            model_management.get_torch_device = lambda: torch.device(self.target_device)
+            # Force clear any cached models that might have wrong device assignments
+            if hasattr(model_management, 'soft_empty_cache'):
+                model_management.soft_empty_cache()
             
-            # Create and execute original preprocessor using its specific function name
-            preprocessor = self.preprocessor_class()
+            # Clear preprocessor instance cache to force fresh loading
+            self._preprocessor_instance = None
+            # Override with consistent device for ENTIRE execution
+            target_torch_device = torch.device(self.target_device)
+            model_management.get_torch_device = lambda: target_torch_device
+            
+            # Ensure all input tensors are on the target device
+            kwargs = self._ensure_tensors_on_device(kwargs, target_torch_device)
+            
+            # Create preprocessor instance if not cached, or reuse existing
+            if self._preprocessor_instance is None:
+                self._preprocessor_instance = self.preprocessor_class()
+            
+            # Execute with consistent device context
             function_name = getattr(self.preprocessor_class, 'FUNCTION', 'execute')
-            result = getattr(preprocessor, function_name)(**kwargs)
+            result = getattr(self._preprocessor_instance, function_name)(**kwargs)
+            
+            # Ensure output tensors are properly placed
+            result = self._ensure_tensors_on_device(result, target_torch_device)
+            
             return result
             
         except Exception as e:
             logger.error(f"Error in MultiGPUPreprocessorWrapper execution: {e}")
+            logger.error(f"Target device: {self.target_device}")
+            logger.error(f"Input tensor devices: {[v.device if isinstance(v, torch.Tensor) else type(v) for v in kwargs.values()]}")
             raise
             
         finally:
@@ -208,4 +242,4 @@ if MidasDepthMapWrapper:
     NODE_CLASS_MAPPINGS["MidasDepthMapWrapper"] = MidasDepthMapWrapper
     NODE_DISPLAY_NAME_MAPPINGS["MidasDepthMapWrapper"] = "Midas Depth (GPU Wrapper)"
 
-logger.info(f"Registered {len(NODE_CLASS_MAPPINGS)} GPU wrapper nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
\ No newline at end of file
+logger.info(f"Registered {len(NODE_CLASS_MAPPINGS)} GPU wrapper nodes: {list(NODE_CLASS_MAPPINGS.keys())}")