digital-garbageandClaude Opus 5.5
109ea2fb4f
Reverts b015ce3f.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
85 lines
3.9 KiB
Python
85 lines
3.9 KiB
Python
"""H3 video decode: spatial tiles, and the X2 Detail VAE's packed output.
|
|
|
|
Ported from v4 (itself from TripleHeadedMonkey/ComfyUI-MiniMaxH3_LatentUpscaler's
|
|
"VAE Decode (fast)"). Two things the stock decode cannot do on H3:
|
|
|
|
* Honour a tile size. `decode_tiled` ignores the tile arguments on H3, and ComfyUI's
|
|
out-of-memory fallback is a generic 3D tiler that ignores H3's blend rules. Here the
|
|
VAE's OWN tile flags are set for the one call and restored.
|
|
* Read an X2 Detail VAE, whose decoder emits packed 12-channel output that has to be
|
|
pixel-shuffled to twice the resolution. Stock decode cannot read it at all.
|
|
|
|
Unvalidated on a GPU, same as v4's own status.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
TILE_MIN = 256 # smaller tiles show a block grid
|
|
SPATIAL_RATIO = 16 # the VAE's spatial compression; the tiler assumes it
|
|
|
|
|
|
def x2_ratio(vae) -> int:
|
|
"""PixelShuffle factor of an H3 video VAE whose decoder emits packed RGB:
|
|
1 for the stock VAE, 2 for the X2 Detail VAE. 1 when it cannot be read."""
|
|
try:
|
|
d = vae.first_stage_model.decoder
|
|
rows = int(d.proj_out.weight.shape[0])
|
|
packed = rows // (int(d.patch_size_t) * int(d.patch_size) ** 2) // 3
|
|
r = math.isqrt(packed)
|
|
return r if r >= 1 and r * r == packed and rows % 3 == 0 else 1
|
|
except Exception: # noqa: BLE001
|
|
return 1
|
|
|
|
|
|
def unpack(out, up):
|
|
"""[B, 3*up*up, T, H, W] packed phases -> [B, 3, T, H*up, W*up]. Layout is R phases,
|
|
then G, then B (what `repeat_interleave` on the pixel statistics matches)."""
|
|
b, c, t, h, w = out.shape
|
|
out = F.pixel_shuffle(out.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w), up)
|
|
return out.reshape(b, t, c // (up * up), h * up, w * up).permute(0, 2, 1, 3, 4)
|
|
|
|
|
|
def tiles(tile_px):
|
|
"""(tile, overlap) in pixels: at least TILE_MIN, a multiple of 16, overlap a
|
|
quarter of the tile (bigger tiles with a fixed 64px overlap show a block grid)."""
|
|
tile = max(TILE_MIN, (int(tile_px) // SPATIAL_RATIO) * SPATIAL_RATIO)
|
|
return tile, max(16, (tile * 64 // 256 // 16) * 16)
|
|
|
|
|
|
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."""
|
|
import comfy.model_management as mm
|
|
inner = vae.first_stage_model
|
|
up = x2_ratio(vae)
|
|
tile, overlap = tiles(tile_px)
|
|
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
|
|
if up > 1:
|
|
inner.decoder.out_channels = 3 * up * up
|
|
with mm.cuda_device_context(vae.device):
|
|
mm.load_models_gpu([vae.patcher],
|
|
memory_required=vae.memory_used_decode(video.shape, vae.vae_dtype),
|
|
force_full_load=vae.disable_offload)
|
|
if up > 1:
|
|
# After the load: a (dynamic) load can restore the checkpoint's 3-channel stats.
|
|
inner.pixel_mean = inner.pixel_mean.repeat_interleave(up * up, dim=1)
|
|
inner.pixel_std = inner.pixel_std.repeat_interleave(up * up, dim=1)
|
|
out = torch.empty(inner.decode_output_shape(video.shape),
|
|
device=mm.intermediate_device(), dtype=vae.vae_output_dtype())
|
|
inner.decode(video.to(device=vae.device, dtype=vae.vae_dtype), output_buffer=out)
|
|
vae.process_output(out)
|
|
finally:
|
|
(inner.tiling, inner.tile_size, inner.tile_overlap_min,
|
|
inner.decoder.out_channels, inner.pixel_mean, inner.pixel_std) = saved
|
|
if up > 1:
|
|
out = unpack(out, up)
|
|
return out.movedim(1, -1)
|