Fix skip_until_exit logic bug breaking injection after first run of m…

…odel
comfyanonymous · Sep 21, 2024 · f28d892 · f28d892
1 parent 5f450d3
commit f28d892
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
@@ -180,7 +180,10 @@ def __enter__(self):
 
     def __exit__(self, *args):
         if self.was_injected:
-            if self.skip_until_exit or not self.model.skip_injection:
+            if self.skip_until_exit:
+                self.model.skip_injection = self.prev_skip_injection
+                self.model.inject_model()
+            elif not self.model.skip_injection:
                 self.model.inject_model()
         self.model.skip_injection = self.prev_skip_injection