Fix FreeU not working when shape is tensor.

comfyanonymous · comfyanonymous · commit b02bcced05dc · 2024-05-23T11:48:04.000-04:00
diff --git a/comfy_extras/nodes_freelunch.py b/comfy_extras/nodes_freelunch.py
@@ -42,7 +42,7 @@ def patch(self, model, b1, b2, s1, s2):
         on_cpu_devices = {}
 
         def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(h.shape[1], None)
+            scale = scale_dict.get(int(h.shape[1]), None)
             if scale is not None:
                 h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0]
                 if hsp.device not in on_cpu_devices:
@@ -81,7 +81,7 @@ def patch(self, model, b1, b2, s1, s2):
         on_cpu_devices = {}
 
         def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(h.shape[1], None)
+            scale = scale_dict.get(int(h.shape[1]), None)
             if scale is not None:
                 hidden_mean = h.mean(1).unsqueeze(1)
                 B = hidden_mean.shape[0]
