Files

93 lines
4.5 KiB
Python

"""Latent to pictures, and to sound where there is any.
Writing the file is FunUI Save Video's job (core's SaveVideo encoded on one
CPU thread); VHS is unnecessary either way.
What core cannot generalise is the decode itself for a model whose latent is not
one tensor. An AV latent carries a video branch and an audio branch that need
different VAEs, and how they are arranged is the model's business. So the same
rule as the empty latent applies: a model's own module may claim the decode, and
otherwise it is the ordinary single-tensor path.
The node takes the MODEL for one reason: so a provider can answer "is this mine"
by IDENTITY rather than by shape. Without it, the H3 module claimed any latent
that happened to have two parts -- correct today, because H3 is the only thing
producing one, and silently wrong the moment a second AV model exists. Deciding
from shape is how a model mismatch turns into an unrelated-looking fault instead
of an error.
"""
from comfy_api.latest import io
from ..._core import log, registry as registry_mod, vae_tiles
CAPABILITY = "decode"
def _plain_decode(vae, latent, tile_size):
return vae_tiles.tiled_decode(vae, latent, tile_size)
class FunUIDecode(io.ComfyNode):
@classmethod
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="FunUIDecode",
display_name="FunUI Decode",
category="FunUI/Output",
description="Decode a latent to images, and to audio when the model has any.",
inputs=[
io.Latent.Input("samples"),
io.Vae.Input("vae", tooltip="The picture VAE."),
io.Model.Input("model", optional=True,
tooltip="The model this latent came from. Needed only when a "
"model has its own way of being decoded."),
io.Vae.Input("audio_vae", optional=True,
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="Decodes in tiles of this many pixels to fit less VRAM. 0 uses the model's own way. For MiniMax H3 keep 0, since bigger tiles show a grid."),
],
outputs=[
io.Image.Output(display_name="images"),
io.Audio.Output(display_name="audio"),
io.String.Output(display_name="status"),
],
)
@classmethod
def execute(cls, samples, vae, model=None, audio_vae=None, tile_size=0) -> io.NodeOutput:
latent = samples["samples"]
vae = getattr(vae, "funui_decode_vae", None) or vae # a decode-only VAE named in the VAE loader
for spec, decode in registry_mod.current().providers(CAPABILITY):
try:
claimed = decode(latent, model=model, vae=vae, audio_vae=audio_vae, tile_size=tile_size)
except Exception as exc: # noqa: BLE001
# Same protocol as the empty latent: returning None means "not my
# model", so getting this far means it WAS and it broke. Falling
# back would decode a nested latent as if it were one tensor.
log.broke(f"{spec.id}.{CAPABILITY}", exc, "decoding this model's latent")
raise RuntimeError(
f"{spec.id} handles this model's decode and failed: "
f"{type(exc).__name__}: {exc}. Refusing to decode it as a plain "
f"latent, which would not be the same picture."
) from exc
if claimed is not None:
images, audio = claimed
return io.NodeOutput(images, audio, f"{spec.id} decoded this")
if getattr(latent, "is_nested", False):
# Nothing claimed it and it is not one tensor: vae.decode would either
# raise somewhere confusing or quietly decode the wrong branch.
raise RuntimeError(
"This latent has more than one part and no installed module knows "
"how to decode it."
+ ("" if model is not None else
" The model input is not wired, so no module could recognise it.")
+ " The module for this model is missing.")
images = _plain_decode(vae, latent, tile_size)
if len(images.shape) == 5: # a batch of clips -> a strip of frames
images = images.reshape(-1, *images.shape[-3:])
return io.NodeOutput(images, None, "decoded as a single latent")