Various improvements to bislerp.

comfy/utils.py

@@ -56,35 +56,42 @@ def bislerp(samples, width, height):
     shape[2] = height
     out1 = torch.empty(shape, dtype=samples.dtype, layout=samples.layout, device=samples.device)
 
-    def algorithm(in1, w1, in2, w2):
+    def algorithm(in1, in2, t):
         dims = in1.shape
-        val = w2
+        val = t
 
         #flatten to batches
         low = in1.reshape(dims[0], -1)
         high = in2.reshape(dims[0], -1)
 
-        low_norm = low/torch.norm(low, dim=1, keepdim=True)
-        high_norm = high/torch.norm(high, dim=1, keepdim=True)
-
-        # in case we divide by zero
-        low_norm[low_norm != low_norm] = 0.0
-        high_norm[high_norm != high_norm] = 0.0
-
-        omega = torch.acos((low_norm*high_norm).sum(1))
+        low_weight = torch.norm(low, dim=1, keepdim=True)
+        low_weight[low_weight == 0] = 0.0000000001
+        low_norm = low/low_weight
+        high_weight = torch.norm(high, dim=1, keepdim=True)
+        high_weight[high_weight == 0] = 0.0000000001
+        high_norm = high/high_weight
+
+        dot_prod = (low_norm*high_norm).sum(1)
+        dot_prod[dot_prod > 0.9995] = 0.9995
+        dot_prod[dot_prod < -0.9995] = -0.9995
+        omega = torch.acos(dot_prod)
         so = torch.sin(omega)
-        res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
+        res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low_norm + (torch.sin(val*omega)/so).unsqueeze(1) * high_norm
+        res *= (low_weight * (1.0-val) + high_weight * val)
         return res.reshape(dims)
 
     for x_dest in range(shape[3]):
         for y_dest in range(shape[2]):
-            y = (y_dest) * height_scale
-            x = (x_dest) * width_scale
+            y = (y_dest + 0.5) * height_scale - 0.5
+            x = (x_dest + 0.5) * width_scale - 0.5
 
             x1 = max(math.floor(x), 0)
             x2 = min(x1 + 1, samples.shape[3] - 1)
+            wx = x - math.floor(x)
+
             y1 = max(math.floor(y), 0)
             y2 = min(y1 + 1, samples.shape[2] - 1)
+            wy = y - math.floor(y)
 
             in1 = samples[:,:,y1,x1]
             in2 = samples[:,:,y1,x2]
@@ -94,13 +101,13 @@ def bislerp(samples, width, height):
             if (x1 == x2) and (y1 == y2):
                 out_value = in1
             elif (x1 == x2):
-                out_value = algorithm(in1, (y2 - y), in3, (y - y1))
+                out_value = algorithm(in1, in3, wy)
             elif (y1 == y2):
-                out_value = algorithm(in1, (x2 - x), in2, (x - x1))
+                out_value = algorithm(in1, in2, wx)
             else:
-                o1 = algorithm(in1, (x2 - x), in2, (x - x1))
-                o2 = algorithm(in3, (x2 - x), in4, (x - x1))
-                out_value = algorithm(o1, (y2 - y), o2, (y - y1))
+                o1 = algorithm(in1, in2, wx)
+                o2 = algorithm(in3, in4, wx)
+                out_value = algorithm(o1, o2, wy)
 
             out1[:,:,y_dest,x_dest] = out_value
     return out1