From b015ce3f87f4b82c6064910af9a1c3cc3a16d350 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:51:11 +0300 Subject: [PATCH] fix: H3 decode with tile size 0 decodes each frame in one piece (core tiled at 256 regardless); out of memory steps down to 512, then 256 tiles Co-Authored-By: Claude Opus 5.5 --- modules/models/minimax_h3/__init__.py | 17 ++++++++--- .../minimax_h3/tests/test_vae_decode.py | 28 +++++++++++++------ modules/models/minimax_h3/vae_decode.py | 10 ++++--- modules/output/decode/nodes.py | 5 ++-- 4 files changed, 42 insertions(+), 18 deletions(-) diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 81edb53e..8c1796cc 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -175,7 +175,8 @@ def _video_images(vae, video_latent, tile_size): """The picture branch. An X2 Detail VAE can only be read by the tiled decoder (its output is packed), so it always goes that way -- and has no stock path to fall back to, so a failure there is raised, naming the VAE. For a stock VAE, tiles are - asked for by `tile_size` and, if the tiled decode fails, said once and decoded plain.""" + asked for by `tile_size` (0 = each frame in one piece) and, if that decode fails, said + once and decoded the stock way.""" x2 = _decode.x2_ratio(vae) if x2 > 1: try: @@ -184,12 +185,20 @@ def _video_images(vae, video_latent, tile_size): raise RuntimeError( f"this VAE is an X2 Detail VAE (x{x2} packed output) and only the tiled decoder " f"can read it, and that failed: {type(exc).__name__}: {exc}") from exc - if tile_size and tile_size > 0: + import comfy.model_management as mm + # 0 = one piece (fastest); out of memory steps down to big tiles, then core's 256. + for size in ([tile_size] if tile_size and tile_size > 0 else [0, 512, 256]): try: - return _decode.decode_fast(vae, video_latent, tile_size) + return _decode.decode_fast(vae, video_latent, size) except Exception as exc: # noqa: BLE001 + if "out of memory" in str(exc).lower(): + mm.soft_empty_cache() + log.once(f"h3_decode_oom:{size}", log.ALERT, "FunPack H3 decode", + f"out of memory decoding {'in one piece' if not size else f'in {size}px tiles'}; trying smaller tiles") + continue log.once(f"h3_decode_tiles:{type(exc).__name__}", log.ALERT, "FunPack H3 decode", - f"the tiled decode failed ({type(exc).__name__}: {exc}); decoded in one piece instead") + f"the fast decode failed ({type(exc).__name__}: {exc}); decoded the stock way instead (slower)") + break return vae.decode(video_latent) diff --git a/modules/models/minimax_h3/tests/test_vae_decode.py b/modules/models/minimax_h3/tests/test_vae_decode.py index 8f7cf346..7a6ba819 100644 --- a/modules/models/minimax_h3/tests/test_vae_decode.py +++ b/modules/models/minimax_h3/tests/test_vae_decode.py @@ -30,6 +30,7 @@ class _Inner: self.pixel_mean, self.pixel_std = torch.arange(3.).view(1, 3, 1, 1, 1), torch.ones(1, 3, 1, 1, 1) self.during = None self.fail = False + self.tried = [] def decode_output_shape(self, shape): b, _, t, h, w = shape @@ -38,8 +39,9 @@ class _Inner: def decode(self, x, output_buffer): self.during = (self.tiling, self.tile_size, self.tile_overlap_min, self.decoder.out_channels, self.pixel_mean.shape[1]) + self.tried.append(self.tile_size if self.tiling else 0) if self.fail: - raise RuntimeError("out of memory") + raise RuntimeError(self.fail if isinstance(self.fail, str) else "CUDA out of memory") output_buffer.copy_(torch.arange(output_buffer.numel(), dtype=output_buffer.dtype).view_as(output_buffer)) @@ -151,10 +153,10 @@ def _provider_decode(vae, tile): return h3._video_images(vae, torch.zeros(1, 4, 2, 8, 8), tile) -def test_tile_size_zero_on_a_stock_vae_is_the_plain_decode(patched): +def test_tile_size_zero_on_a_stock_vae_decodes_each_frame_in_one_piece(patched): vae = _Vae(12) - out = _provider_decode(vae, 0) - assert vae.plain_calls == 1 and vae.first_stage_model.during is None and out.shape == (1, 2, 4, 4, 3) + _provider_decode(vae, 0) + assert vae.plain_calls == 0 and vae.first_stage_model.during[0] is False # core's own 256px tiling is off def test_a_stock_vae_with_tiles_goes_through_the_tiled_decoder(patched): @@ -163,14 +165,24 @@ def test_a_stock_vae_with_tiles_goes_through_the_tiled_decoder(patched): assert vae.plain_calls == 0 and vae.first_stage_model.during[:3] == (True, 512, 128) -def test_a_failing_tiled_decode_falls_back_to_plain_and_says_so_once(patched): +def test_out_of_memory_in_one_piece_steps_down_to_512_then_256_then_stock(patched): from core import log log._reset() vae = _Vae(12) vae.first_stage_model.fail = True - out = _provider_decode(vae, 512) - assert vae.plain_calls == 1 and out.shape == (1, 2, 4, 4, 3) - assert any("decoded in one piece" in r["message"] for r in log.history()) + out = _provider_decode(vae, 0) + assert vae.first_stage_model.tried == [0, 512, 256] and vae.plain_calls == 1 and out.shape == (1, 2, 4, 4, 3) + assert any("out of memory decoding in one piece" in r["message"] for r in log.history()) + + +def test_any_other_failure_goes_straight_to_the_stock_decode_and_says_so(patched): + from core import log + log._reset() + vae = _Vae(12) + vae.first_stage_model.fail = "shape mismatch" + _provider_decode(vae, 512) + assert vae.first_stage_model.tried == [512] and vae.plain_calls == 1 + assert any("stock way" in r["message"] for r in log.history()) def test_an_x2_vae_never_falls_back_it_names_itself(patched): diff --git a/modules/models/minimax_h3/vae_decode.py b/modules/models/minimax_h3/vae_decode.py index 82ea7317..d9421015 100644 --- a/modules/models/minimax_h3/vae_decode.py +++ b/modules/models/minimax_h3/vae_decode.py @@ -51,17 +51,19 @@ def tiles(tile_px): def decode_fast(vae, video, tile_px): """Decode an H3 video latent in `tile_px`-pixel spatial tiles straight through H3's - decoder. Returns IMAGES (B, T, H, W, C) on the intermediate device. The VAE's flags - are restored even when the decode raises.""" + decoder; `tile_px` 0 = each frame in one piece (core's own default is 256px tiles with + 64px overlaps: ~1.8x the pixels in many small pieces, the slow part). Returns IMAGES + (B, T, H, W, C) on the intermediate device. The VAE's flags are restored even when the + decode raises.""" import comfy.model_management as mm inner = vae.first_stage_model up = x2_ratio(vae) - tile, overlap = tiles(tile_px) + tile, overlap = tiles(tile_px) if tile_px and tile_px > 0 else (inner.tile_size, inner.tile_overlap_min) saved = (inner.tiling, inner.tile_size, inner.tile_overlap_min, inner.decoder.out_channels, inner.pixel_mean, inner.pixel_std) vae.throw_exception_if_invalid() try: - inner.tiling, inner.tile_size, inner.tile_overlap_min = True, tile, overlap + inner.tiling, inner.tile_size, inner.tile_overlap_min = bool(tile_px and tile_px > 0), tile, overlap if up > 1: inner.decoder.out_channels = 3 * up * up with mm.cuda_device_context(vae.device): diff --git a/modules/output/decode/nodes.py b/modules/output/decode/nodes.py index 053eea42..06a35980 100644 --- a/modules/output/decode/nodes.py +++ b/modules/output/decode/nodes.py @@ -49,8 +49,9 @@ class FunPackDecode(io.ComfyNode): tooltip="Only for models that generate sound alongside the video."), io.Int.Input("tile_size", default=0, min=0, max=4096, step=64, optional=True, tooltip="Decode in tiles of this many pixels, to fit less VRAM. 0 = in one " - "piece. Bigger = faster and more VRAM. An X2 Detail VAE is always " - "tiled (256 at least)."), + "piece, the fastest (MiniMax H3: falls back to 512, then 256 tiles " + "if memory runs out). Bigger = faster and more VRAM. An X2 Detail " + "VAE is always tiled (256 at least)."), ], outputs=[ io.Image.Output(display_name="images"),