fix: ensure patched torch graph is always synced on inference errors (#129)

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

This commit ensures that the patched torch graph remains synchronized
even when an error occurs during inference, preventing potential
inconsistencies and adds logging for controlnet tensor cloning errors
---------

Co-authored-by: Elite <john@eliteencoder.net>

Author: rickstaa

nodes/tensor_utils/prestartup_script.py

@@ -28,26 +28,29 @@ 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)}")
+                raise 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()