Fix MiniMax H3 tiled VAE decode producing a grid of mismatched tiles

Each spatial tile decoded by MiniMaxH3VideoVAE.tiled_decode was
normalized to its own local -1..1 rotary position range instead of its
true offset within the full latent grid, since tiling defaults to on
(tile_size=256) at any resolution above that. The ViT3D decoder then
saw every tile as if it were the whole frame, so adjacent tiles came
out with different learned appearance and visible seams (#15548).

create_token_ids now takes the full latent grid size and the tile's
offset within it, and tiled_decode passes them through so tile
position ids match a single-shot decode of the whole frame.

comfy/ldm/minimax/vae.py

@@ -166,11 +166,17 @@ class EncoderFCN3D(nn.Module):
 
 # ViT3D decoder
 
-def create_token_ids(patch_dims, device, dtype):
+def create_token_ids(patch_dims, device, dtype, full_dims=None, offset=None):
+    # full_dims/offset let a spatially tiled chunk compute positions relative to the
+    # full latent grid instead of restarting at -1..1 for every tile (see tiled_decode).
+    if full_dims is None:
+        full_dims = patch_dims
+    if offset is None:
+        offset = (0,) * len(patch_dims)
     coords_list = []
-    for dim_size in patch_dims:
-        coords = torch.arange(0.5, dim_size, dtype=dtype, device=device)
-        coords = coords / dim_size
+    for dim_size, full_size, off in zip(patch_dims, full_dims, offset):
+        coords = torch.arange(0.5, dim_size, dtype=dtype, device=device) + off
+        coords = coords / full_size
         coords = 2.0 * coords - 1.0
         coords_list.append(coords)
     coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)
@@ -282,7 +288,7 @@ class ViT3DDecoder(nn.Module):
         self.norm_out = ops.LayerNorm(dim, elementwise_affine=True, eps=eps)
         self.proj_out = ops.Linear(dim, out_channels * patch_size_t * patch_size * patch_size)
 
-    def forward(self, x):
+    def forward(self, x, full_dims=None, offset=None):
         B, C, latent_T, latent_H, latent_W = x.shape
 
         h = self.x_embedder(x.flatten(2).transpose(1, 2))  # [B, T*H*W, C]
@@ -292,7 +298,8 @@ class ViT3DDecoder(nn.Module):
 
         h = torch.cat([h, comfy.ops.cast_to_input(self.register_tokens, h).expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1)
 
-        img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype).expand(B, -1, -1)
+        img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype,
+                                    full_dims=full_dims, offset=offset).expand(B, -1, -1)
         suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)
         img_ids = torch.cat([img_ids, suffix_ids], dim=1)
 
@@ -389,8 +396,8 @@ class MiniMaxH3VideoVAE(nn.Module):
     def _encode_moments(self, x):
         return self.quant_conv(self.encoder(x))
 
-    def _decode_pixels(self, z):
-        return self.decoder(self.post_quant_conv(z))
+    def _decode_pixels(self, z, full_dims=None, offset=None):
+        return self.decoder(self.post_quant_conv(z), full_dims=full_dims, offset=offset)
 
     def _normalize_pixels(self, x):
         return x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
@@ -504,6 +511,7 @@ class MiniMaxH3VideoVAE(nn.Module):
 
     def tiled_decode(self, z):
         height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
+        full_dims = (z.shape[-3], z.shape[-2], z.shape[-1])
         y_idx, y_len, y_overlap = self.split_tiles(height)
         x_idx, x_len, x_overlap = self.split_tiles(width)
 
@@ -518,7 +526,7 @@ class MiniMaxH3VideoVAE(nn.Module):
             out_x = 0
             for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
                 zj, zw = j_pos // self.vae_ratio, j_len // self.vae_ratio
-                tile = self._decode_pixels(z[..., zi:zi + zl, zj:zj + zw])
+                tile = self._decode_pixels(z[..., zi:zi + zl, zj:zj + zw], full_dims=full_dims, offset=(0, zi, zj))
                 if i < len(y_idx) - 1:
                     new_tails.append(tile[..., -y_overlap[i]:, :].clone())
                 next_left_tail = tile[..., :, -x_overlap[j]:].clone() if j < len(x_idx) - 1 else None

tests-unit/comfy_test/test_minimax_h3_vae_tiled_decode.py

@@ -0,0 +1,34 @@
+import torch
+
+from comfy.cli_args import args as cli_args
+
+if not torch.cuda.is_available():
+    cli_args.cpu = True
+
+from comfy.ldm.minimax.vae import create_token_ids  # noqa: E402
+
+
+def test_create_token_ids_defaults_match_full_range():
+    # full_dims/offset omitted must reproduce the original (pre-tiling-fix) behavior.
+    baseline = create_token_ids((1, 4, 4), device="cpu", dtype=torch.float32)
+    same = create_token_ids((1, 4, 4), device="cpu", dtype=torch.float32, full_dims=(1, 4, 4), offset=(0, 0, 0))
+    torch.testing.assert_close(baseline, same)
+
+
+def test_create_token_ids_tile_offset_matches_global_slice():
+    # A spatial tile decoded at its true offset within the full latent grid must get the
+    # same position ids as the corresponding slice of a single-shot full-grid decode.
+    # Without full_dims/offset, every tile is normalized to its own local -1..1 range
+    # regardless of where it sits in the frame, which is the root cause of ComfyUI/#15548
+    # (MiniMax H3 tiled decode producing a grid of independently textured tiles).
+    full = create_token_ids((1, 8, 8), device="cpu", dtype=torch.float32).view(1, 8, 8, 3)
+
+    top_left = create_token_ids(
+        (1, 4, 4), device="cpu", dtype=torch.float32, full_dims=(1, 8, 8), offset=(0, 0, 0)
+    ).view(1, 4, 4, 3)
+    bottom_right = create_token_ids(
+        (1, 4, 4), device="cpu", dtype=torch.float32, full_dims=(1, 8, 8), offset=(0, 4, 4)
+    ).view(1, 4, 4, 3)
+
+    torch.testing.assert_close(top_left, full[:, 0:4, 0:4, :])
+    torch.testing.assert_close(bottom_right, full[:, 4:8, 4:8, :])