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.
This commit is contained in:
@@ -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())}")
|
||||
logger.info(f"Registered {len(NODE_CLASS_MAPPINGS)} GPU wrapper nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
|
||||
Reference in New Issue
Block a user