fix(api): pad partial latents when working with tiles

ssube · Jun 6, 2023 · 6b1ce37 · 6b1ce37
1 parent 2a16041
commit 6b1ce37
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/api/onnx_web/diffusers/utils.py b/api/onnx_web/diffusers/utils.py
@@ -272,7 +272,15 @@ def get_tile_latents(
     xt = x + t
     yt = y + t
 
-    return full_latents[:, :, y:yt, x:xt]
+    tile_latents = full_latents[:, :, y:yt, x:xt]
+
+    if tile_latents.shape[2] < t or tile_latents.shape[3] < t:
+        px = t - tile_latents.shape[3]
+        py = t - tile_latents.shape[2]
+
+        tile_latents = np.pad(tile_latents, ((0, 0), (0, 0), (0, py), (0, px)), mode="reflect")
+
+    return tile_latents
 
 
 def get_scaled_latents(