Limited the U-Net restoration message to only display when it happens

dvruette · Aug 26, 2023 · eafc798 · eafc798
1 parent 2f207ff
commit eafc798
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/scripts/fabric.py b/scripts/fabric.py
@@ -461,7 +461,6 @@ def process(self, p, *args):
             print("[FABRIC] Skipping U-Net forward pass patching")
 
     def postprocess(self, p, processed, *args):
-        print("[FABRIC] Restoring original U-Net forward pass")
         unpatch_unet_forward_pass(p.sd_model.model.diffusion_model)
 
         images = processed.images[processed.index_of_first_image:]

diff --git a/scripts/patching.py b/scripts/patching.py
@@ -212,5 +212,6 @@ def patched_attn1_forward(attn1, idx, x, context=None, **kwargs):
 
 def unpatch_unet_forward_pass(unet):
     if hasattr(unet, "_fabric_old_forward"):
+        print("[FABRIC] Restoring original U-Net forward pass")
         unet.forward = unet._fabric_old_forward
         del unet._fabric_old_forward