From b02bcced05dcc2d09a0358eff6393c5641485878 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 23 May 2024 11:47:43 -0400 Subject: [PATCH] Fix FreeU not working when shape is tensor. --- comfy_extras/nodes_freelunch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_freelunch.py b/comfy_extras/nodes_freelunch.py index 6f1d87bf..c5ebcf26 100644 --- a/comfy_extras/nodes_freelunch.py +++ b/comfy_extras/nodes_freelunch.py @@ -42,7 +42,7 @@ class FreeU: 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 @@ class FreeU_V2: 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]