output: decode, which is the last thing FunPack needs to own
Core already turns images into a video and writes the file, and handles containers, codecs and muxing properly. Writing our own would be a worse copy of tested code, so there is no FunPack save node and no reason for VHS either. What core cannot generalise is decoding a latent that is not one tensor: an AV latent has two branches wanting different VAEs, and their arrangement is the model's business. Same rule as the empty latent -- a model module may claim it, and H3 now does. A nested latent nobody claims is refused rather than handed to vae.decode, which would raise somewhere confusing or quietly decode one branch as the picture. A missing audio VAE is named, because returning no audio looks exactly like a model that makes none. With this the pipeline closes end to end using FunPack's eight nodes and core's samplers and video nodes -- verified by loading both together. That matters more than the node count: until now nothing FunPack produced could be measured, only reasoned about. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Opus 5
Claude Opus 5
parent
749b1b5db9
commit
134cee78b7
@@ -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",
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Decode. A node-only module."""
|
||||
|
||||
from .nodes import FunPackDecode
|
||||
|
||||
ID = "output_decode"
|
||||
TITLE = "Decode"
|
||||
STAGE = "post"
|
||||
STATUS = "proven"
|
||||
|
||||
NODES = [FunPackDecode]
|
||||
@@ -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")
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user