Make HyperTile deterministic

Author: asagi4

comfy_extras/nodes_hypertile.py

@@ -2,18 +2,18 @@
 
 import math
 from einops import rearrange
-import random
+# Use torch rng for consistency across generations
+from torch import randint
 
-def random_divisor(value: int, min_value: int, /, max_options: int = 1, counter = 0) -> int:
+def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
     min_value = min(min_value, value)
 
     # All big divisors of value (inclusive)
     divisors = [i for i in range(min_value, value + 1) if value % i == 0]
 
     ns = [value // i for i in divisors[:max_options]]  # has at least 1 element
 
-    random.seed(counter)
-    idx = random.randint(0, len(ns) - 1)
+    idx = randint(low=0, high=len(ns) - 1, size=(1,)).item()
 
     return ns[idx]
 
@@ -42,7 +42,6 @@ def patch(self, model, tile_size, swap_size, max_depth, scale_depth):
 
         latent_tile_size = max(32, tile_size) // 8
         self.temp = None
-        self.counter = 1
 
         def hypertile_in(q, k, v, extra_options):
             if q.shape[-1] in apply_to:
@@ -53,10 +52,8 @@ def hypertile_in(q, k, v, extra_options):
                 h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio))
 
                 factor = 2**((q.shape[-1] // model_channels) - 1) if scale_depth else 1
-                nh = random_divisor(h, latent_tile_size * factor, swap_size, self.counter)
-                self.counter += 1
-                nw = random_divisor(w, latent_tile_size * factor, swap_size, self.counter)
-                self.counter += 1
+                nh = random_divisor(h, latent_tile_size * factor, swap_size)
+                nw = random_divisor(w, latent_tile_size * factor, swap_size)
 
                 if nh * nw > 1:
                     q = rearrange(q, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw)