192 lines
9.6 KiB
Python
192 lines
9.6 KiB
Python
"""PromptForge MiniMax H3 VAE Decode: the H3 video decoder with an out-of-memory ladder.
|
|
|
|
The stock decode gives up on an OOM by switching to a generic tiler that ignores H3's blend rules. This node
|
|
calls H3's own decoder instead and, when memory runs out, steps down a ladder that keeps the same quality:
|
|
|
|
256 px tiles (the size H3's decoder was built for) -> 85 / 34 / 17 frames per pass
|
|
|
|
Whole-frame or larger tiles look fried and banded with H3, so they are never chosen automatically.
|
|
The rung that worked for a given latent size is remembered, so the next run does not repeat the failed
|
|
attempts. Output is written into one preallocated buffer (no list-then-concatenate), non-finite pixels are
|
|
repaired and reported, and a model that has a 2x packed output projection (like the Wan 2x VAE trick) is
|
|
pixel-shuffled a few frames at a time.
|
|
"""
|
|
import gc
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import comfy.model_management as model_management
|
|
import comfy.utils
|
|
|
|
CONTEXT_GROUPS = 2
|
|
# (spatial tile; frame groups per pass, 0 = whole clip). One group = 17 frames, 5 latent steps.
|
|
LADDER = ((256, 0), (256, 5), (256, 2), (256, 1))
|
|
SHUFFLE_FRAMES = 8
|
|
_worked = {}
|
|
|
|
|
|
def _inner(vae):
|
|
inner = getattr(vae, "first_stage_model", None)
|
|
if not all(hasattr(inner, name) for name in ("tiling", "tile_size", "tile_overlap_min", "decode_output_shape")):
|
|
raise ValueError("[PromptForge] MiniMax H3 VAE Decode needs the MiniMax H3 video VAE (not the audio VAE).")
|
|
return inner
|
|
|
|
|
|
def _video(samples):
|
|
latent = samples["samples"]
|
|
video = latent.unbind()[0] if latent.is_nested else latent
|
|
if video.ndim != 5:
|
|
raise ValueError(f"[PromptForge] expected an H3 video latent [B, C, T, H, W], got {tuple(video.shape)}.")
|
|
return video
|
|
|
|
|
|
def _packing(inner):
|
|
"""PixelShuffle factor of a decoder whose projection emits packed RGB sub-pixels (1 for the stock H3 decoder)."""
|
|
decoder = inner.decoder
|
|
channels = decoder.proj_out.weight.shape[0] // (decoder.patch_size_t * decoder.patch_size ** 2)
|
|
factor = math.isqrt(channels // 3)
|
|
if channels % 3 or factor * factor * 3 != channels:
|
|
raise ValueError(f"[PromptForge] the H3 decoder emits {channels} channels per pixel, which is not a packed RGB layout.")
|
|
return channels, factor
|
|
|
|
|
|
def _expand(stat, channels):
|
|
return stat if stat.shape[1] == channels else stat.repeat_interleave(channels // 3, dim=1)
|
|
|
|
|
|
def _shuffle(pixels, factor):
|
|
"""[B, C, T, H, W] packed channels -> [B, T, H * factor, W * factor, 3], a few frames at a time."""
|
|
batch, channels, frames, height, width = pixels.shape
|
|
out = torch.empty(batch, frames, height * factor, width * factor, channels // factor ** 2, device=pixels.device, dtype=pixels.dtype)
|
|
for start in range(0, frames, SHUFFLE_FRAMES):
|
|
part = pixels[:, :, start:start + SHUFFLE_FRAMES]
|
|
count = part.shape[2]
|
|
part = F.pixel_shuffle(part.movedim(1, 2).reshape(batch * count, channels, height, width), factor)
|
|
out[:, start:start + count] = part.reshape(batch, count, -1, height * factor, width * factor).movedim(2, -1)
|
|
return out
|
|
|
|
|
|
def _decode(vae, video, device, factor):
|
|
"""One decode pass of `video` with the tiling currently set on the model; returns [B, T, H, W, 3] on `device`."""
|
|
vae.throw_exception_if_invalid()
|
|
inner = vae.first_stage_model
|
|
with model_management.cuda_device_context(vae.device):
|
|
model_management.load_models_gpu([vae.patcher], memory_required=vae.memory_used_decode(video.shape, vae.vae_dtype), force_full_load=vae.disable_offload)
|
|
if factor > 1: # a dynamic load can restore the 3-channel statistics, so expand them after it
|
|
inner.pixel_mean = _expand(inner.pixel_mean, inner.decoder.out_channels)
|
|
inner.pixel_std = _expand(inner.pixel_std, inner.decoder.out_channels)
|
|
pixels = torch.empty(inner.decode_output_shape(video.shape), device=device, dtype=vae.vae_output_dtype())
|
|
inner.decode(video.to(device=vae.device, dtype=vae.vae_dtype), output_buffer=pixels)
|
|
vae.process_output(pixels)
|
|
return _shuffle(pixels, factor) if factor > 1 else pixels.movedim(1, -1)
|
|
|
|
|
|
def _windows(latent_steps, groups):
|
|
"""(first latent step, end latent step, frames to drop from the front) for each pass over a 2 + 5n latent."""
|
|
total = (latent_steps - 2) // 5
|
|
done = 0
|
|
while done < total:
|
|
core = min(groups, total - done)
|
|
context = min(CONTEXT_GROUPS, done)
|
|
yield 5 * (done - context), 2 + 5 * (done + core), (5 + 17 * context) if done else 0
|
|
done += core
|
|
|
|
|
|
def _run(vae, video, device, factor, rung):
|
|
tile, groups = rung
|
|
inner = vae.first_stage_model
|
|
inner.tiling = tile > 0
|
|
if tile:
|
|
inner.tile_size, inner.tile_overlap_min = tile, max(64, tile // 4)
|
|
if not groups or video.shape[2] < 2 or (video.shape[2] - 2) % 5 or (video.shape[2] - 2) // 5 <= groups:
|
|
return _decode(vae, video, device, factor)
|
|
|
|
windows = list(_windows(video.shape[2], groups))
|
|
progress = comfy.utils.ProgressBar(len(windows))
|
|
out = None
|
|
position = 0
|
|
for start, end, drop in windows:
|
|
model_management.throw_exception_if_processing_interrupted()
|
|
piece = _decode(vae, video[:, :, start:end], device, factor)[:, drop:]
|
|
if out is None:
|
|
frames = 5 + 17 * ((video.shape[2] - 2) // 5)
|
|
out = torch.empty(piece.shape[0], frames, *piece.shape[2:], device=piece.device, dtype=piece.dtype)
|
|
out[:, position:position + piece.shape[1]] = piece
|
|
position += piece.shape[1]
|
|
progress.update(1)
|
|
del piece
|
|
return out[:, :position]
|
|
|
|
|
|
class PROMPTFORGE_MiniMaxDecode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE", {"tooltip": "The MiniMax H3 video VAE."}),
|
|
"tile_size": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 16, "tooltip": "0 = auto: 256 px tiles, the size H3's decoder was built for. Larger tiles or the whole frame decode faster but look fried and banded. Overlap is a quarter of the tile."}),
|
|
"frames_per_pass": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 17, "tooltip": "0 = auto: the whole clip, then fewer frames per pass on out-of-memory. A number decodes the clip in passes of about that many frames, with causal context."}),
|
|
"output_device": (["cpu", "gpu"], {"default": "cpu", "tooltip": "Where the finished frames are kept. gpu skips the copy to RAM but needs VRAM for the whole clip."}),
|
|
"retry_on_oom": ("BOOLEAN", {"default": True, "tooltip": "On out-of-memory, step down to smaller tiles and shorter passes instead of failing."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("images", "INFO")
|
|
FUNCTION = "decode"
|
|
CATEGORY = "🔮 PromptForge/Video"
|
|
DESCRIPTION = "MiniMax H3 video decode with 256px tiles that steps down to shorter passes on out-of-memory instead of failing, and remembers what worked."
|
|
|
|
def decode(self, samples, vae, tile_size, frames_per_pass, output_device, retry_on_oom):
|
|
inner = _inner(vae)
|
|
video = _video(samples)
|
|
channels, factor = _packing(inner)
|
|
device = model_management.get_torch_device() if output_device == "gpu" else model_management.intermediate_device()
|
|
key = tuple(video.shape)
|
|
|
|
if tile_size or frames_per_pass:
|
|
forced = (max(256, tile_size // 16 * 16), max(1, round(frames_per_pass / 17)) if frames_per_pass else 0)
|
|
unlimited = lambda value: value or 10 ** 9
|
|
safer = [rung for rung in LADDER if unlimited(rung[0]) < unlimited(forced[0]) or rung[0] == forced[0] and unlimited(rung[1]) < unlimited(forced[1])]
|
|
ladder = [forced] + (safer if retry_on_oom else [])
|
|
else:
|
|
ladder = list(LADDER[_worked.get(key, 0):]) if retry_on_oom else [LADDER[_worked.get(key, 0)]]
|
|
|
|
saved = (inner.tiling, inner.tile_size, inner.tile_overlap_min, inner.decoder.out_channels, inner.pixel_mean, inner.pixel_std)
|
|
inner.decoder.out_channels = channels
|
|
failures = []
|
|
try:
|
|
for index, rung in enumerate(ladder):
|
|
try:
|
|
images = _run(vae, video, device, factor, rung)
|
|
except model_management.OOM_EXCEPTION:
|
|
failures.append(rung)
|
|
gc.collect()
|
|
model_management.soft_empty_cache()
|
|
if index == len(ladder) - 1:
|
|
raise
|
|
continue
|
|
break
|
|
finally:
|
|
inner.tiling, inner.tile_size, inner.tile_overlap_min, inner.decoder.out_channels, inner.pixel_mean, inner.pixel_std = saved
|
|
|
|
if not tile_size and not frames_per_pass and rung in LADDER:
|
|
_worked[key] = LADDER.index(rung)
|
|
repaired = False
|
|
for start in range(0, images.shape[1], 16):
|
|
part = images[:, start:start + 16]
|
|
if not torch.isfinite(part).all():
|
|
torch.nan_to_num_(part, nan=0.0, posinf=1.0, neginf=0.0)
|
|
repaired = True
|
|
|
|
how = f"{rung[0]}px tiles" + (f", {17 * rung[1]} frames per pass" if rung[1] else "")
|
|
info = f"{images.shape[1]} frames {images.shape[3]}x{images.shape[2]}, {how}" + (f" (after {len(failures)} out-of-memory retries)" if failures else "") + (", non-finite pixels were repaired" if repaired else "")
|
|
return (images.reshape(-1, *images.shape[-3:]), info)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_MiniMaxDecode": PROMPTFORGE_MiniMaxDecode}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_MiniMaxDecode": "PromptForge MiniMax H3 VAE Decode"}
|