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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-10-07 14:51:11 +03:00
co-authored by Claude Opus 5.5
parent dc67fc8274
commit b015ce3f87
4 changed files with 42 additions and 18 deletions
+13 -4
View File
@@ -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)
@@ -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):
+6 -4
View File
@@ -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):
+3 -2
View File
@@ -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"),