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:
co-authored by
Claude Opus 5.5
parent
dc67fc8274
commit
b015ce3f87
@@ -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):
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"),
|
||||
|
||||
Reference in New Issue
Block a user