298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-20 03:13:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make bislerp work on GPU.

Browse Source
This commit is contained in:
comfyanonymous 2023-11-14 11:38:36 -05:00
parent 420beeeb05
commit c962884a5c
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
comfy/utils.py
View File

@ -307,13 +307,13 @@ def bislerp(samples, width, height):
res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1] res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]
return res return res
def generate_bilinear_data(length_old, length_new): def generate_bilinear_data(length_old, length_new, device):
coords_1 = torch.arange(length_old).reshape((1,1,1,-1)).to(torch.float32) coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))
coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear") coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")
ratios = coords_1 - coords_1.floor() ratios = coords_1 - coords_1.floor()
coords_1 = coords_1.to(torch.int64) coords_1 = coords_1.to(torch.int64)
coords_2 = torch.arange(length_old).reshape((1,1,1,-1)).to(torch.float32) + 1 coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1
coords_2[:,:,:,-1] -= 1 coords_2[:,:,:,-1] -= 1
coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear") coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")
coords_2 = coords_2.to(torch.int64) coords_2 = coords_2.to(torch.int64)
@ -323,7 +323,7 @@ def bislerp(samples, width, height):
h_new, w_new = (height, width) h_new, w_new = (height, width)
#linear w #linear w
ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new) ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)
coords_1 = coords_1.expand((n, c, h, -1)) coords_1 = coords_1.expand((n, c, h, -1))
coords_2 = coords_2.expand((n, c, h, -1)) coords_2 = coords_2.expand((n, c, h, -1))
ratios = ratios.expand((n, 1, h, -1)) ratios = ratios.expand((n, 1, h, -1))
@ -336,7 +336,7 @@ def bislerp(samples, width, height):
result = result.reshape(n, h, w_new, c).movedim(-1, 1) result = result.reshape(n, h, w_new, c).movedim(-1, 1)
#linear h #linear h
ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new) ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)
coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new)) coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new)) coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new)) ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))