diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 25fae92e..b89cd94e 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -48,6 +48,7 @@ FROZEN = { ], "FunPackLoadModifiers": ["model", "settings"], "FunPackModifierSettings": ["settings"], + "FunPackDecode": ["samples", "vae", "audio_vae"], "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size"], "FunPackDiffusionModelLoader": [ "model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation", diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index f4446f94..48948efc 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -59,6 +59,37 @@ def traits(model): return found +def decode(latent, vae, audio_vae=None): + """H3's two branches, each through the VAE that understands it. + + Returns None when the latent is not nested, so this never speaks for a model + it does not own -- the branch order is H3's own arrangement, and reading it + wrongly would decode noise as a picture rather than fail. + """ + if not getattr(latent, "is_nested", False): + return None + + parts = latent.unbind() + if len(parts) != 2: + raise RuntimeError( + f"expected a video and an audio branch, got {len(parts)} parts") + video_latent, audio_latent = parts + + images = vae.decode(video_latent) + if len(images.shape) == 5: + images = images.reshape(-1, *images.shape[-3:]) + + if audio_vae is None: + # Silence would look like a decoded soundtrack. Say what is missing. + raise RuntimeError( + "this model generates sound and no audio VAE is wired into the " + "decode, so there is nothing to turn the audio branch into") + + from comfy_extras.nodes_audio import vae_decode_audio + audio = vae_decode_audio(audio_vae, {"samples": audio_latent}) + return images, audio + + def empty_latent(model, width, height, length, batch_size=1): """H3's joint video+audio latent, or None if this is not an H3 model.""" if not is_h3(model): @@ -77,4 +108,4 @@ def empty_latent(model, width, height, length, batch_size=1): TRAITS = traits -PROVIDES = {"empty_latent": empty_latent} +PROVIDES = {"empty_latent": empty_latent, "decode": decode} diff --git a/modules/output/__init__.py b/modules/output/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modules/output/decode/__init__.py b/modules/output/decode/__init__.py new file mode 100644 index 00000000..e404be41 --- /dev/null +++ b/modules/output/decode/__init__.py @@ -0,0 +1,10 @@ +"""Decode. A node-only module.""" + +from .nodes import FunPackDecode + +ID = "output_decode" +TITLE = "Decode" +STAGE = "post" +STATUS = "proven" + +NODES = [FunPackDecode] diff --git a/modules/output/decode/nodes.py b/modules/output/decode/nodes.py new file mode 100644 index 00000000..89773253 --- /dev/null +++ b/modules/output/decode/nodes.py @@ -0,0 +1,75 @@ +"""Latent to pictures, and to sound where there is any. + +This is the only part of getting a result out that FunPack needs to own. Core +already turns images into a video (`CreateVideo`) and writes the file +(`SaveVideo`, `SaveWEBM`, `SaveImage`), and those handle containers, codecs and +audio muxing properly -- writing our own would be duplication with a worse +version of somebody else's tested code. 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. +""" + +from comfy_api.latest import io + +from ..._core import log, registry as registry_mod + +CAPABILITY = "decode" + + +class FunPackDecode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="FunPackDecode", + display_name="FunPack Decode", + category="FunPack/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.Vae.Input("audio_vae", optional=True, + tooltip="Only for models that generate sound alongside the video."), + ], + 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, audio_vae=None) -> io.NodeOutput: + latent = samples["samples"] + + for spec, decode in registry_mod.current().providers(CAPABILITY): + try: + claimed = decode(latent, vae=vae, audio_vae=audio_vae) + 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.failed(f"{spec.id}.{CAPABILITY}", exc) + 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. The module for this model is missing.") + + images = vae.decode(latent) + 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") diff --git a/modules/output/tests/test_decode.py b/modules/output/tests/test_decode.py new file mode 100644 index 00000000..8e6cb409 --- /dev/null +++ b/modules/output/tests/test_decode.py @@ -0,0 +1,117 @@ +"""Decoding: the ordinary path, and the one a model has to claim.""" + +import pytest + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """Imports comfy.""" + + +class _Vae: + def __init__(self, out=None): + self.seen = [] + self._out = out + + def decode(self, latent): + import torch + self.seen.append(latent) + return self._out if self._out is not None else torch.zeros(4, 8, 8, 3) + + +def test_a_single_tensor_latent_decodes_normally(): + import torch + from modules.output.decode.nodes import FunPackDecode + + vae = _Vae() + out = FunPackDecode.execute({"samples": torch.zeros(1, 4, 8, 8)}, vae) + images, audio, status = out.result + + assert audio is None and "single latent" in status + assert vae.seen and images.shape[-1] == 3 + + +def test_a_batch_of_clips_becomes_one_strip_of_frames(): + """A 5-D decode is [B, T, H, W, C]; downstream wants frames.""" + import torch + from modules.output.decode.nodes import FunPackDecode + + vae = _Vae(out=torch.zeros(2, 5, 8, 8, 3)) + images, _audio, _status = FunPackDecode.execute({"samples": torch.zeros(1, 4, 8, 8)}, vae).result + assert images.shape == (10, 8, 8, 3) + + +def test_a_nested_latent_nobody_claims_is_refused_not_decoded_as_one_tensor(monkeypatch): + """vae.decode on a NestedTensor would either raise somewhere confusing or + quietly decode one branch as if it were the picture. + + H3 is installed and claims any nested latent, so proving this path needs its + claim taken away -- which is what a machine without that module would look + like. + """ + import torch + from comfy.nested_tensor import NestedTensor + from core import registry as registry_mod + from modules.output.decode.nodes import FunPackDecode + + for spec in registry_mod.current().specs.values(): + if "decode" in spec.provides: + monkeypatch.delitem(spec.provides, "decode") + + nested = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 16))) + with pytest.raises(RuntimeError, match="more than one part"): + FunPackDecode.execute({"samples": nested}, _Vae()) + + +def test_h3_claims_a_nested_latent_and_uses_both_vaes(monkeypatch): + import torch + from comfy.nested_tensor import NestedTensor + from modules.models import minimax_h3 + + video_vae = _Vae(out=torch.zeros(1, 3, 8, 8, 3)) + audio_vae = _Vae() + + monkeypatch.setattr("comfy_extras.nodes_audio.vae_decode_audio", + lambda vae, samples: {"waveform": "decoded", "sample_rate": 32000}) + + nested = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 16))) + images, audio = minimax_h3.decode(nested, vae=video_vae, audio_vae=audio_vae) + + assert images.shape == (3, 8, 8, 3) + assert audio["waveform"] == "decoded" + + +def test_h3_does_not_speak_for_a_plain_latent(): + import torch + from modules.models import minimax_h3 + assert minimax_h3.decode(torch.zeros(1, 4, 8, 8), vae=_Vae()) is None + + +def test_a_missing_audio_vae_is_named_rather_than_silently_silent(): + """Returning no audio would look exactly like a model that makes none.""" + import torch + from comfy.nested_tensor import NestedTensor + from modules.models import minimax_h3 + + nested = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 16))) + with pytest.raises(RuntimeError, match="audio VAE"): + minimax_h3.decode(nested, vae=_Vae(out=torch.zeros(1, 3, 8, 8, 3)), audio_vae=None) + + +def test_a_claiming_module_that_breaks_stops_the_decode(monkeypatch): + import torch + from comfy.nested_tensor import NestedTensor + from core import registry as registry_mod + from modules.output.decode.nodes import FunPackDecode + + def broken(latent, vae, audio_vae=None): + if not getattr(latent, "is_nested", False): + return None + raise RuntimeError("branch order changed") + + spec = registry_mod.current().specs["model_minimax_h3"] + monkeypatch.setitem(spec.provides, "decode", broken) + + nested = NestedTensor((torch.zeros(1, 24, 2, 8, 8), torch.zeros(1, 32, 2, 16))) + with pytest.raises(RuntimeError, match="Refusing to decode"): + FunPackDecode.execute({"samples": nested}, _Vae())