wuyingwen/comfyui
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Fix FreeU not working when shape is tensor.

Browse Source
This commit is contained in:
comfyanonymous 2024-05-23 11:47:43 -04:00
parent 6507a9c716
commit b02bcced05
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
comfy_extras/nodes_freelunch.py
View File

@ -42,7 +42,7 @@ class FreeU:
on_cpu_devices = {} on_cpu_devices = {}
def output_block_patch(h, hsp, transformer_options): 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: if scale is not None:
h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0] h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0]
if hsp.device not in on_cpu_devices: if hsp.device not in on_cpu_devices:
@ -81,7 +81,7 @@ class FreeU_V2:
on_cpu_devices = {} on_cpu_devices = {}
def output_block_patch(h, hsp, transformer_options): 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: if scale is not None:
hidden_mean = h.mean(1).unsqueeze(1) hidden_mean = h.mean(1).unsqueeze(1)
B = hidden_mean.shape[0] B = hidden_mean.shape[0]