fix for SAG with Kohya HRFix/ Deep Shrink (#5546)

now works with arbitrary downscale factors

Author: DenOfEquity

comfy_extras/nodes_sag.py

@@ -57,12 +57,17 @@ def create_blur_map(x0, attn, sigma=3.0, threshold=1.0):
     attn = attn.reshape(b, -1, hw1, hw2)
     # Global Average Pool
     mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold
-    ratio = 2**(math.ceil(math.sqrt(lh * lw / hw1)) - 1).bit_length()
-    mid_shape = [math.ceil(lh / ratio), math.ceil(lw / ratio)]
+
+    f = float(lh) / float(lw)
+    fh = f ** 0.5
+    fw = (1/f) ** 0.5
+    S = mask.size(1) ** 0.5
+    w = int(0.5 + S * fw)
+    h = int(0.5 + S * fh)
 
     # Reshape
     mask = (
-        mask.reshape(b, *mid_shape)
+        mask.reshape(b, h, w)
         .unsqueeze(1)
         .type(attn.dtype)
     )