nodes/Brioch-ComfyUI-MiniMaxH3-Preview
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
Brioch-ComfyUI-MiniMaxH3-Pr…/test_audio_preview.py
T

487 lines
23 KiB
Python
Raw Permalink Blame History

"""Offline checks for the audio preview — no server, no GPU, no model.
<ComfyUI>/venv/bin/python test_audio_preview.py [--write-clip]
Needs torch, because what it protects is tensor plumbing: which of the packed streams gets
picked up, the scale the sampler carries the audio at, how a window in seconds becomes latent
frames, and whether the bytes handed to the browser are a real audio file. The ComfyUI modules
the node imports are stubbed below — none of them are involved in the parts under test, and
stubbing them is what keeps this runnable without a server.
The fake VAEs below take their channel count, sample rate and samples-per-latent-frame as
constructor arguments, so the checks can prove those three come from the VAE and the decode
rather than from constants — a VAE variant at another rate is read correctly instead of quietly
decoding nonsense.
Two checks here exist because of bugs that shipped:
* **The stream index.** Video and audio differ by one index in a packed latent, and getting it
wrong does not raise — it decodes silence, or noise, and reads as a model problem. The stream
is now matched by the VAE's channel count, and a pack with audio first proves it.
* **The audio scale.** `ModelSamplingAV` carries the audio stream at `shift / audio_shift`, so
a decode that does not divide it out returns loud noise at every step while the video stream
looks perfect and points at nothing.
"""
import importlib.util
import io
import logging
import os
import sys
import types
import wave
HERE = os.path.dirname(os.path.abspath(__file__))
try:
import torch
except ImportError:
print("SKIP test_audio_preview.py needs torch — run it with the python that runs "
"ComfyUI:\n <ComfyUI>/venv/bin/python test_audio_preview.py")
sys.exit(0)
# ---------------------------------------------------------------- stub ComfyUI
def _module(name, **attrs):
"""Register a stub module — and hang it off its parent.
`import comfy.utils` finds the entry in sys.modules, but `comfy.utils.unpack_latents` then
reads an *attribute* of `comfy`, which only the real import machinery sets.
"""
m = types.ModuleType(name)
for k, v in attrs.items():
setattr(m, k, v)
sys.modules[name] = m
if "." in name:
parent, _, leaf = name.rpartition(".")
setattr(sys.modules[parent], leaf, m)
return m
def _unpack_latents(packed, shapes):
"""Stand-in for comfy.utils.unpack_latents: split a flat pack back into its streams."""
out, offset = [], 0
for shape in shapes:
n = 1
for dim in shape:
n *= dim
out.append(packed[..., offset:offset + n].reshape((packed.shape[0],) + tuple(shape)))
offset += n
return out
def _vae_decode_audio(vae, samples):
"""Stand-in for comfy_extras.nodes_audio.vae_decode_audio, same contract."""
latent = samples["samples"]
if getattr(latent, "is_nested", False):
latent = latent.unbind()[-1]
return {"waveform": vae.decode(latent),
"sample_rate": getattr(vae, "audio_sample_rate", 32000)}
class _Routes:
"""Collects the routes the module registers, so the test can assert one was."""
def __init__(self):
self.registered = []
def get(self, path):
self.registered.append(path)
return lambda fn: fn
_ROUTES = _Routes()
_module("comfy")
_module("comfy.utils", unpack_latents=_unpack_latents)
_module("comfy.model_management",
OOM_EXCEPTION=getattr(torch.cuda, "OutOfMemoryError", MemoryError))
_module("comfy.patcher_extension",
WrappersMP=types.SimpleNamespace(OUTER_SAMPLE="outer_sample"),
add_wrapper_with_key=lambda *a, **k: None,
get_all_wrappers=lambda *a, **k: [])
_module("server", PromptServer=types.SimpleNamespace(
instance=types.SimpleNamespace(routes=_ROUTES, send_sync=lambda *a, **k: None)))
_module("aiohttp", web=types.SimpleNamespace(Response=lambda **kw: kw))
_module("comfy_extras")
_module("comfy_extras.nodes_audio", vae_decode_audio=_vae_decode_audio)
_module("comfy_api")
_io_stub = types.SimpleNamespace(
ComfyNode=type("ComfyNode", (), {}),
Schema=lambda **kw: kw, NodeOutput=lambda *a, **k: None,
Hidden=types.SimpleNamespace(unique_id="unique_id"))
for _name in ("Model", "Vae", "Int", "Float", "Boolean", "Combo", "Latent"):
setattr(_io_stub, _name, types.SimpleNamespace(Input=lambda *a, **k: None,
Output=lambda *a, **k: None))
_module("comfy_api.latest", io=_io_stub)
spec = importlib.util.spec_from_file_location(
"minimax_h3_audio_preview", os.path.join(HERE, "minimax_h3_audio_preview.py"))
node = importlib.util.module_from_spec(spec)
sys.modules["minimax_h3_audio_preview"] = node
spec.loader.exec_module(node)
_results = []
def check(name, got, want):
_results.append((got == want, name, got, want))
# H3's two streams at the shapes core builds them: video [B,24,T,h,w], audio [B,32,2,T].
# 32 is H3's audio channel count, and appears here as *the fake VAE's* rather than as something
# the node itself knows.
CHANNELS = 32
VIDEO_SHAPE = (24, 7, 6, 8)
AUDIO_SHAPE = (CHANNELS, 2, 40) # 40 latent frames = one second at 40 fps
VIDEO = torch.zeros((1,) + VIDEO_SHAPE)
AUDIO = torch.ones((1,) + AUDIO_SHAPE) # ones, so a mix-up shows up as silence
PACKED = torch.cat([VIDEO.reshape(1, -1), AUDIO.reshape(1, -1)], dim=-1)
SHAPES = [VIDEO_SHAPE, AUDIO_SHAPE]
# The same pair packed the other way round, to prove the stream is found by matching the VAE's
# channel count rather than by trusting audio to be last.
AUDIO_FIRST = torch.cat([AUDIO.reshape(1, -1), VIDEO.reshape(1, -1)], dim=-1)
AUDIO_FIRST_SHAPES = [AUDIO_SHAPE, VIDEO_SHAPE]
class _Nested:
"""What the sampler hands the callback when it has not been flattened."""
is_nested = True
def __init__(self, *streams):
self.tensors = list(streams)
class _UnbindOnly:
"""The same thing on a build that only exposes unbind()."""
is_nested = True
def __init__(self, *streams):
self._streams = list(streams)
def unbind(self):
return list(self._streams)
class _FakeAudioVAE:
"""Shaped like the real one: [B,C,2,T] -> [B,2,T*per_frame].
`per_frame` and `audio_sample_rate` are what set the latent frame rate the node has to
discover — 800 samples per frame at 32 kHz is H3's 40 fps. Records the latent it was handed,
so a check can confirm the sampler's audio_scale was divided out before the decode rather
than trusting the call to have happened.
"""
def __init__(self, channels=CHANNELS, sample_rate=32000, per_frame=800):
self.latent_channels = channels
self.audio_sample_rate = sample_rate
self.per_frame = per_frame
self.seen = None
def decode(self, z):
self.seen = z
return torch.zeros(z.shape[0], 2, z.shape[-1] * self.per_frame)
class _BrokenVAE:
latent_channels = CHANNELS
audio_sample_rate = 32000
def decode(self, z):
raise ValueError("decoder exploded")
# ---------------------------------------------------------------- stream picking
check("the channel count comes from the VAE",
node.vae_latent_channels(_FakeAudioVAE(channels=7)), 7)
check("a VAE that will not say gets no number",
node.vae_latent_channels(types.SimpleNamespace()), None)
check("nor does a nonsense one",
node.vae_latent_channels(types.SimpleNamespace(latent_channels=0)), None)
check("nested latent: the audio stream is the one with the VAE's channels",
tuple(node.audio_stream_of(_Nested(VIDEO, AUDIO), None, CHANNELS).shape),
tuple(AUDIO.shape))
check("nested latent without .tensors: unbind() is used instead",
tuple(node.audio_stream_of(_UnbindOnly(VIDEO, AUDIO), None, CHANNELS).shape),
tuple(AUDIO.shape))
check("flat pack: audio is unpacked and matched",
tuple(node.audio_stream_of(PACKED, SHAPES, CHANNELS).shape), tuple(AUDIO.shape))
check("flat pack: the audio found is the audio packed",
bool(torch.equal(node.audio_stream_of(PACKED, SHAPES, CHANNELS), AUDIO)), True)
# The point of matching by channel count rather than by index: a model that packs audio first
# is found anyway, where `[1]` or `[-1]` would hand back the video stream and decode nonsense.
check("audio packed FIRST is still found, by its channel count",
bool(torch.equal(node.audio_stream_of(AUDIO_FIRST, AUDIO_FIRST_SHAPES, CHANNELS), AUDIO)),
True)
check("audio nested first is still found",
bool(torch.equal(node.audio_stream_of(_Nested(AUDIO, VIDEO), None, CHANNELS), AUDIO)),
True)
# With nothing to match against, last-stream is the only assumption left — and the one every
# AV pack core builds satisfies.
check("with no channel count, the last stream is assumed",
bool(torch.equal(node.audio_stream_of(PACKED, SHAPES), AUDIO)), True)
check("a bare audio latent passes through",
tuple(node.audio_stream_of(AUDIO, None, CHANNELS).shape), tuple(AUDIO.shape))
check("a bare 3D audio latent passes through too",
tuple(node.audio_stream_of(torch.zeros(1, CHANNELS, 40), None, CHANNELS).shape),
(1, CHANNELS, 40))
check("a video-only 5D latent has no audio",
node.audio_stream_of(VIDEO, None, CHANNELS), None)
check("a 4D image latent is not mistaken for audio",
node.audio_stream_of(torch.zeros(1, 16, 64, 64), None, CHANNELS), None)
check("a bare latent whose channels are not the VAE's is refused",
node.audio_stream_of(torch.zeros(1, 8, 2, 40), None, CHANNELS), None)
check("no latent, no audio", node.audio_stream_of(None, None, CHANNELS), None)
check("a flat pack with no shapes cannot be unpacked",
node.audio_stream_of(PACKED, None, CHANNELS), None)
# ---------------------------------------------------------------- the audio scale
def guider(scale=None, says=None):
"""A stand-in guider. `scale` is what model_sampling carries, `says` what the model's
accessor reports — core's returns 1.0 whenever its latent_shapes is unset."""
sampling = types.SimpleNamespace(audio_scale=1.0 if scale is None else scale)
model = types.SimpleNamespace(model_sampling=sampling)
if says is not None:
model.audio_scale = lambda: says
return types.SimpleNamespace(model_patcher=types.SimpleNamespace(model=model))
check("the model's own accessor is the authority",
node.audio_scale_of(guider(4.0, says=12.0 / 3.0), 2), 4.0)
check("an accessor answering 1.0 on a two-stream pack does not end it",
node.audio_scale_of(guider(4.0, says=1.0), 2), 4.0)
check("model_sampling is read when there is no accessor",
node.audio_scale_of(guider(4.0), 2), 4.0)
check("no audio shift means no rescaling", node.audio_scale_of(guider(), 2), 1.0)
check("a single-stream latent is never rescaled", node.audio_scale_of(guider(4.0), 1), 1.0)
check("a model that knows nothing about it is left alone",
node.audio_scale_of(types.SimpleNamespace(), 2), 1.0)
check("a nonsense scale is ignored rather than dividing by zero",
node.audio_scale_of(guider(0.0, says=0.0), 2), 1.0)
vae = _FakeAudioVAE()
node.decode_audio(vae, AUDIO * 4.0, audio_scale=4.0)
check("the scale is divided out before the decode",
bool(torch.allclose(vae.seen, AUDIO)), True)
node.decode_audio(vae, AUDIO, audio_scale=1.0)
check("scale 1.0 leaves the latent untouched",
bool(torch.allclose(vae.seen, AUDIO)), True)
# ---------------------------------------------------------------- the latent frame rate
# Measured from the decode itself — frames in, samples out — so `window_seconds` means seconds
# on a model whose audio latent does not run at H3's 40 fps.
check("800 samples per frame at 32 kHz is 40 fps",
node.measured_latent_fps(120, 120 * 800, 32000), 40.0)
check("1600 samples per frame at 32 kHz is 20 fps",
node.measured_latent_fps(50, 50 * 1600, 32000), 20.0)
check("a different sample rate is accounted for",
node.measured_latent_fps(10, 10 * 960, 48000), 50.0)
check("no samples means nothing was learnt", node.measured_latent_fps(10, 0, 32000, 40.0), 40.0)
check("no frames either", node.measured_latent_fps(0, 8000, 32000, 40.0), 40.0)
check("an absurd answer is discarded rather than believed",
node.measured_latent_fps(1, 1, 32000, 40.0), 40.0)
check("the seed is only a seed", node.SEED_LATENT_FPS, 40.0)
# ---------------------------------------------------------------- window maths
check("a 0.5s window of a 1s latent is 20 frames at 40 fps",
node.audio_latent_window(AUDIO, 0.5, 40.0).shape[-1], 20)
check("the same 0.5s is 10 frames at 20 fps",
node.audio_latent_window(AUDIO, 0.5, 20.0).shape[-1], 10)
check("the window is taken from the head",
bool(torch.equal(node.audio_latent_window(AUDIO, 0.5, 40.0), AUDIO[..., :20])), True)
check("window 0 keeps the whole clip", node.audio_latent_window(AUDIO, 0, 40.0).shape[-1], 40)
check("a window longer than the clip keeps the whole clip",
node.audio_latent_window(AUDIO, 99, 40.0).shape[-1], 40)
check("a tiny window still yields one frame",
node.audio_latent_window(AUDIO, 0.001, 40.0).shape[-1], 1)
check("an unknown rate cannot cut a window, so nothing is cut",
node.audio_latent_window(AUDIO, 0.5, 0.0).shape[-1], 40)
# ---------------------------------------------------------------- clip + envelope
check("[B,C,L] loses its batch dim", tuple(node.to_clip(torch.zeros(1, 2, 800)).shape), (2, 800))
# Not upmixed: a one-channel WAV plays everywhere, and a mono model should not be made to lie
# about having two channels.
check("mono stays mono", tuple(node.to_clip(torch.zeros(1, 800)).shape), (1, 800))
check("more than two channels are cut to two",
tuple(node.to_clip(torch.zeros(1, 4, 800)).shape), (2, 800))
check("the level is left alone here, not clipped",
float(node.to_clip(torch.full((1, 2, 8), 4.0)).max()), 4.0)
# Clipping an over-range clip is what made noise sound like noise and drew the envelope as a
# solid rectangle; scaling keeps the shape so the waveform still says something.
check("an over-range clip is scaled to fit",
float(node.normalise(torch.full((2, 8), 4.0), 4.0).max()), 1.0)
check("scaling preserves the shape",
[round(float(v), 3) for v in node.normalise(torch.tensor([[2.0, -4.0]]), 4.0)[0]],
[0.5, -1.0])
check("a clip already in range is untouched",
float(node.normalise(torch.full((2, 8), 0.25), 0.25).max()), 0.25)
spiked = torch.zeros(2, 8000)
spiked[0, 100] = 1.0 # one spike, left channel, first bucket
env = node.envelope(spiked, buckets=8)
check("one row per channel", len(env), 2)
check("bucket count honoured", len(env[0]), 8)
check("the spike lands in the first bucket", env[0][0], 1.0)
check("silence reads as silence", max(env[1]), 0.0)
check("an empty clip has no envelope", node.envelope(torch.zeros(2, 0)), [])
check("fewer samples than buckets still works", len(node.envelope(torch.zeros(2, 3), 8)[0]), 3)
# ---------------------------------------------------------------- what the browser gets
# 0.5s of a 32 kHz stereo tone, loud enough that a silent output would show up as one.
tone = torch.sin(torch.arange(0, 16000, dtype=torch.float32) * 0.05).repeat(2, 1) * 0.5
raw = node.encode_wav(tone, 32000)
# One format, deterministically, so these checks can verify the actual bytes that reach the
# browser rather than trusting a codec to have produced something playable.
check("the clip is served from a registered route",
node.AUDIO_ROUTE in _ROUTES.registered, True)
check("WAV magic", raw[:4], b"RIFF")
check("the header declares WAVE", raw[8:12], b"WAVE")
with wave.open(io.BytesIO(raw)) as w:
check("WAV is stereo", w.getnchannels(), 2)
check("WAV is 16-bit", w.getsampwidth(), 2)
check("WAV keeps the sample rate", w.getframerate(), 32000)
check("WAV keeps every sample", w.getnframes(), 16000)
frames = w.readframes(w.getnframes())
check("every sample is present in the data chunk", len(frames), 16000 * 2 * 2)
check("the tone is not silence", max(frames) > 0, True)
# `--write-clip` drops the encoded clip next to this file: ground truth for "encoder or
# browser?" — if the file plays in a media player, the encoder is fine.
if "--write-clip" in sys.argv:
path = os.path.join(HERE, "audio_preview_check.wav")
with open(path, "wb") as f:
f.write(raw)
print("wrote %s — %d bytes, a 0.5s 32 kHz stereo tone.\n"
" If that plays in a media player, the encoder is fine.\n" % (path, len(raw)))
# ---------------------------------------------------------------- the gate
def wrapper(vae=None, **kw):
opts = dict(node_id="1", audio_vae=vae if vae is not None else _FakeAudioVAE(),
window_seconds=3.0, start_at_percent=50, every_n_steps=1, max_overhead=15)
opts.update(kw)
return node._AudioOuterSampleWrapper(**opts)
def state(**kw):
s = {"off": False, "warned": False, "logged": False, "sent": 0, "seq": 0,
"cost": 0.0, "finished": 0.0, "clip": 0.0, "steps": 0, "due": 0, "scale": None,
"fps": node.SEED_LATENT_FPS}
s.update(kw)
return s
on = wrapper()
check("no audio before the gate", on._due(4, 20, state()), False)
check("audio once the gate is passed", on._due(9, 20, state()), True)
check("audio on the last step", on._due(19, 20, state()), True)
check("a disabled run stays silent", on._due(19, 20, state(off=True)), False)
check("0% starts from the first step", wrapper(start_at_percent=0)._due(0, 20, state()), True)
check("100% waits for the last step",
wrapper(start_at_percent=100)._due(18, 20, state()), False)
check("every_n_steps skips the steps in between",
wrapper(every_n_steps=4)._due(11, 20, state()), False)
check("every_n_steps keeps the ones it lands on",
wrapper(every_n_steps=4)._due(12, 20, state()), True)
# ---------------------------------------------------------------- the throttle
check("nothing to wait for on the first update", on._skip(100.0, state()), False)
check("a clip still playing holds the next update",
on._skip(100.0, state(finished=99.0, clip=3.0)), True)
check("once it has played through, the next may go",
on._skip(103.5, state(finished=100.0, clip=3.0)), False)
check("an expensive decode is spaced by the overhead budget",
on._skip(105.0, state(finished=100.0, cost=1.0)), True) # 15% -> ~5.7s gap
check("overhead 0 means only the play-through wait",
wrapper(max_overhead=0)._skip(101.0, state(finished=100.0, cost=99.0)), False)
# ---------------------------------------------------------------- the payload
s = state()
payload = wrapper(window_seconds=0.5)._payload(PACKED, SHAPES, s)
check("a payload is produced", sorted(payload or {}),
["audio_mime", "envelope", "kb", "seconds", "url"])
check("the window decides the duration, not the latent", payload["seconds"], 0.5)
check("the envelope has both channels", len(payload["envelope"]), 2)
check("the url points at the route, for this node, with a sequence number",
payload["url"], node.AUDIO_ROUTE + "?node_id=1&seq=1")
check("the bytes are parked for the route to serve",
node._AUDIO_CLIPS["1"][1][:4], b"RIFF")
check("the parked clip is the whole clip",
len(node._AUDIO_CLIPS["1"][1]), 44 + 16000 * 2 * 2)
check("the payload's size matches what was parked",
payload["kb"], len(node._AUDIO_CLIPS["1"][1]) // 1024)
check("a second update gets a fresh url so nothing is answered from cache",
wrapper(window_seconds=0.5)._payload(PACKED, SHAPES, s)["url"],
node.AUDIO_ROUTE + "?node_id=1&seq=2")
check("window 0 decodes the whole second",
wrapper(window_seconds=0)._payload(PACKED, SHAPES, state())["seconds"], 1.0)
# ------------------------------------------- a VAE that is not H3's, end to end
# 1600 samples per frame at 32 kHz is a 20 fps audio latent, half H3's. The first window is cut
# at the seed rate and so comes out twice as long as asked; the decode measures the real rate,
# and every window after it is right. Nothing in the node was told the rate.
slow = _FakeAudioVAE(per_frame=1600)
slow_latent = torch.ones(1, CHANNELS, 2, 200) # 10s at 20 fps
w = wrapper(vae=slow, window_seconds=1.0)
s = state()
check("the first window is cut at the seed rate, so it runs long",
w._payload(slow_latent, None, s)["seconds"], 2.0)
check("and the decode measured the real frame rate", s["fps"], 20.0)
check("so the next window is the length actually asked for",
w._payload(slow_latent, None, s)["seconds"], 1.0)
# A 48 kHz mono decoder: rate, channel count and frame rate all differ from H3's, and all three
# come from the VAE and the decode rather than from a constant.
mono = _FakeAudioVAE(channels=8, sample_rate=48000, per_frame=960) # 50 fps
mono.decode = lambda z: torch.zeros(z.shape[0], 1, z.shape[-1] * 960)
s = state()
mono_payload = wrapper(vae=mono, window_seconds=0.5)._payload(
torch.ones(1, 8, 2, 400), None, s)
check("an 8-channel audio latent is decoded too", mono_payload is not None, True)
check("its measured frame rate is its own", s["fps"], 50.0)
check("one lane in the envelope for a mono clip", len(mono_payload["envelope"]), 1)
with wave.open(io.BytesIO(node._AUDIO_CLIPS["1"][1])) as _w:
check("and a mono WAV reaches the browser", _w.getnchannels(), 1)
check("at the VAE's own sample rate", _w.getframerate(), 48000)
# The wrong VAE is now caught by the mismatch rather than by a hard-coded channel count — and
# it says so at warning level, which a passing test should not print.
logging.disable(logging.CRITICAL)
try:
s = state()
check("a VAE whose channels match nothing in the latent decodes nothing",
wrapper(vae=_FakeAudioVAE(channels=13))._payload(PACKED, SHAPES, s), None)
check("and that is a wiring mistake, not a failure to switch off for", s["off"], False)
finally:
logging.disable(logging.NOTSET)
# ---------------------------------------------------------------- failures stay silent
# The next check makes the decoder throw on purpose, and the code under test logs that with
# exc_info — a real traceback in the output of a passing test reads as a failure, so mute it.
logging.disable(logging.CRITICAL)
try:
s = state()
check("a decode failure returns no payload rather than raising",
wrapper(vae=_BrokenVAE())._payload(PACKED, SHAPES, s), None)
check("and it turns audio off for the rest of the run", s["off"], True)
s = state()
check("a video-only latent yields no payload", on._payload(VIDEO, None, s), None)
check("and that is not treated as a failure", s["off"], False)
finally:
logging.disable(logging.NOTSET)
# ---------------------------------------------------------------- report
failed = [r for r in _results if not r[0]]
for ok, name, got, want in _results:
if not ok:
print("FAIL %s\n got: %r\n want: %r" % (name, got, want))
print("\n%d checks, %d passed, %d failed" %
(len(_results), len(_results) - len(failed), len(failed)))
sys.exit(1 if failed else 0)
Reference in New Issue View Git Blame Copy Permalink