diff --git a/comfy/ldm/minimax/vae.py b/comfy/ldm/minimax/vae.py index e1c146ec4..9fc32d371 100644 --- a/comfy/ldm/minimax/vae.py +++ b/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 diff --git a/tests-unit/comfy_test/test_minimax_h3_vae_tiled_decode.py b/tests-unit/comfy_test/test_minimax_h3_vae_tiled_decode.py new file mode 100644 index 000000000..1be936bd0 --- /dev/null +++ b/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, :])