fix: ensure patched torch graph is always synced on inference errors

nodes/tensor_utils/prestartup_script.py

@@ -28,26 +28,28 @@ def wrapped_control_merge(self, control, control_prev, output_dtype):
                 }
 
             # Get result from original merge function
-            result = original_control_merge(self, control, control_prev, output_dtype)
-
-            # Clone all output tensors
-            result = {
-                k: [t.clone() if t is not None else None for t in v]
-                for k, v in result.items()
-            }
-
-            # Mark CUDA graph step at end
-            if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
-                torch.compiler.cudagraph_mark_step_begin()
-                torch.cuda.synchronize()
-
-            return result
-
+            try:
+                result = original_control_merge(self, control, control_prev, output_dtype)
+
+                # Clone all output tensors
+                result = {
+                    k: [t.clone() if t is not None else None for t in v]
+                    for k, v in result.items()
+                }
+                return result
+            except Exception as e:
+                print(f"Error: Failed to clone ControlNet during inference: {str(e)}")
+            finally:
+                # Mark CUDA graph step at end
+                if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
+                    torch.compiler.cudagraph_mark_step_begin()
+                    torch.cuda.synchronize()
+
         # Apply the patch
         ControlBase.control_merge = wrapped_control_merge
         print("Successfully patched ControlNet for torch.compile() compatibility")
     except Exception as e:
         print(f"Warning: Failed to patch ControlNet: {str(e)}")
 
 # Apply patch when module is imported
-patch_controlnet_for_stream() 
+patch_controlnet_for_stream()