danrh
8f23c95546
Add measured-audio, timed-lyrics, H3 lane planning, identity relay, upload-song routing, model-specific camera contracts, LTX long-horizon planning, workflow migrations, the validated V6 example, regression coverage, and an Unreleased changelog.
508 lines
20 KiB
Python
508 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
import uuid
|
|
from unittest import mock
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
COMFY_ROOT = ROOT.parents[1]
|
|
|
|
|
|
def load_timed_module():
|
|
module_name = f"diffusiongemma_timed_lyrics_fixture_{uuid.uuid4().hex}"
|
|
spec = importlib.util.spec_from_file_location(module_name, ROOT / "timed_lyrics_nodes.py")
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("could not load timed lyrics nodes")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def load_package():
|
|
if str(COMFY_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
package_name = f"diffusiongemma_timed_package_fixture_{uuid.uuid4().hex}"
|
|
spec = importlib.util.spec_from_file_location(
|
|
package_name,
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("could not load DiffusionGemma Prompt Builder")
|
|
package = importlib.util.module_from_spec(spec)
|
|
sys.modules[package_name] = package
|
|
spec.loader.exec_module(package)
|
|
return package
|
|
|
|
|
|
def audio_fixture(
|
|
*,
|
|
duration: float = 12.0,
|
|
sample_rate: int = 1_000,
|
|
vocal_sections: tuple[tuple[float, float], ...] = ((2.0, 5.0),),
|
|
):
|
|
samples = int(round(duration * sample_rate))
|
|
time = torch.arange(samples, dtype=torch.float32) / sample_rate
|
|
accompaniment = (
|
|
0.06 * torch.sin(2.0 * torch.pi * 73.0 * time)
|
|
+ 0.025 * torch.sin(2.0 * torch.pi * 117.0 * time)
|
|
)
|
|
vocals = torch.zeros_like(time)
|
|
for start, end in vocal_sections:
|
|
mask = (time >= start) & (time < end)
|
|
# A deterministic voice-like signal which is independent of the mix bed.
|
|
phrase = (
|
|
0.15 * torch.sin(2.0 * torch.pi * 181.0 * time)
|
|
+ 0.05 * torch.sin(2.0 * torch.pi * 263.0 * time)
|
|
)
|
|
vocals = torch.where(mask, phrase, vocals)
|
|
final = {
|
|
"waveform": (accompaniment + vocals).reshape(1, 1, -1).contiguous(),
|
|
"sample_rate": sample_rate,
|
|
}
|
|
stem = {
|
|
"waveform": vocals.reshape(1, 1, -1).contiguous(),
|
|
"sample_rate": sample_rate,
|
|
}
|
|
return final, stem
|
|
|
|
|
|
class FakeTranscriber:
|
|
def __init__(self, responses=None, exception: Exception | None = None):
|
|
self.responses = list(responses or [])
|
|
self.exception = exception
|
|
self.calls: list[dict[str, float | int | str]] = []
|
|
|
|
def __call__(
|
|
self,
|
|
audio,
|
|
language,
|
|
return_timestamps,
|
|
chunk_offset_seconds,
|
|
absolute_song_start_seconds,
|
|
):
|
|
self.calls.append(
|
|
{
|
|
"duration": audio["waveform"].shape[-1] / audio["sample_rate"],
|
|
"language": language,
|
|
"return_timestamps": int(bool(return_timestamps)),
|
|
"chunk_offset_seconds": chunk_offset_seconds,
|
|
"absolute_song_start_seconds": absolute_song_start_seconds,
|
|
}
|
|
)
|
|
if self.exception is not None:
|
|
raise self.exception
|
|
if self.responses:
|
|
return self.responses.pop(0)
|
|
return {
|
|
"text": "we light the night",
|
|
"chunks": [
|
|
{"text": "we light the night", "timestamp": [0.25, 1.75]},
|
|
],
|
|
}
|
|
|
|
|
|
class FakeModel:
|
|
def __init__(self):
|
|
self.config = SimpleNamespace(_name_or_path="fake/whisper-tiny")
|
|
self.moves: list[str] = []
|
|
|
|
def to(self, device):
|
|
self.moves.append(str(device))
|
|
return self
|
|
|
|
|
|
def fake_pipeline(transcriber: FakeTranscriber, *, with_model: bool = False):
|
|
result = {"model_id": "fake/whisper-tiny", "transcribe": transcriber}
|
|
if with_model:
|
|
result["model"] = FakeModel()
|
|
return result
|
|
|
|
|
|
class TimedLyricsAnalyzerTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
cls.timed = load_timed_module()
|
|
|
|
def setUp(self) -> None:
|
|
self.timed._ANALYSIS_CACHE.clear()
|
|
|
|
def run_lyrics(
|
|
self,
|
|
final,
|
|
stem,
|
|
pipeline,
|
|
*,
|
|
lyrics="[Verse 1]\nWe light the night\nAnd hold the sky",
|
|
song_duration=None,
|
|
excerpt_start=0.0,
|
|
excerpt_duration=None,
|
|
confidence=0.55,
|
|
):
|
|
duration = final["waveform"].shape[-1] / final["sample_rate"]
|
|
return self.timed.DiffusionGemmaTimedLyricsAnalyzer().analyze(
|
|
final,
|
|
"Lyrics + lip sync",
|
|
lyrics,
|
|
self.timed._waveform_sha256(final),
|
|
duration if song_duration is None else song_duration,
|
|
excerpt_start,
|
|
duration if excerpt_duration is None else excerpt_duration,
|
|
"en",
|
|
confidence,
|
|
stem,
|
|
pipeline,
|
|
)
|
|
|
|
def test_input_schema_and_lazy_surface(self) -> None:
|
|
node = self.timed.DiffusionGemmaTimedLyricsAnalyzer
|
|
inputs = node.INPUT_TYPES()
|
|
required = inputs["required"]
|
|
optional = inputs["optional"]
|
|
self.assertEqual(required["performance_mode"][0], "STRING")
|
|
self.assertTrue(required["performance_mode"][1]["forceInput"])
|
|
self.assertTrue(required["lyrics"][1]["forceInput"])
|
|
self.assertTrue(required["lyrics"][1]["multiline"])
|
|
self.assertEqual(required["language"][1]["default"], "en")
|
|
self.assertAlmostEqual(required["minimum_alignment_confidence"][1]["default"], 0.55)
|
|
self.assertTrue(optional["vocal_stem"][1]["lazy"])
|
|
self.assertTrue(optional["whisper_pipeline"][1]["lazy"])
|
|
self.assertEqual(
|
|
node.check_lazy_status(
|
|
performance_mode="Lyrics + lip sync",
|
|
vocal_stem=None,
|
|
whisper_pipeline=None,
|
|
),
|
|
["vocal_stem", "whisper_pipeline"],
|
|
)
|
|
self.assertEqual(
|
|
node.check_lazy_status(
|
|
performance_mode="Lyrics + lip sync",
|
|
vocal_stem=object(),
|
|
whisper_pipeline=None,
|
|
),
|
|
["whisper_pipeline"],
|
|
)
|
|
self.assertEqual(
|
|
node.check_lazy_status(
|
|
performance_mode="Natural / audio-led sync",
|
|
vocal_stem=None,
|
|
whisper_pipeline=None,
|
|
),
|
|
[],
|
|
)
|
|
|
|
def test_natural_and_dance_never_evaluate_expensive_lazy_inputs(self) -> None:
|
|
final, _stem = audio_fixture()
|
|
|
|
class ExplodesOnAccess:
|
|
def __getattribute__(self, _name):
|
|
raise AssertionError("lazy input was evaluated")
|
|
|
|
node = self.timed.DiffusionGemmaTimedLyricsAnalyzer()
|
|
for mode in ("Natural / audio-led sync", "Dance / music sync"):
|
|
outputs = node.analyze(
|
|
final,
|
|
mode,
|
|
"We light the night",
|
|
"not-needed",
|
|
12.0,
|
|
0.0,
|
|
12.0,
|
|
"en",
|
|
0.55,
|
|
ExplodesOnAccess(),
|
|
ExplodesOnAccess(),
|
|
)
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "not_required")
|
|
self.assertFalse(report["timing_ready"])
|
|
self.assertTrue(outputs[2])
|
|
self.assertEqual(report["events"], [])
|
|
|
|
def test_exact_full_mix_separator_fallback_is_rejected_safely(self) -> None:
|
|
final, _stem = audio_fixture()
|
|
transcriber = FakeTranscriber()
|
|
outputs = self.run_lyrics(
|
|
final,
|
|
{"waveform": final["waveform"].clone(), "sample_rate": final["sample_rate"]},
|
|
fake_pipeline(transcriber),
|
|
)
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertFalse(report["timing_ready"])
|
|
self.assertTrue(outputs[2])
|
|
self.assertIn("vocal_stem_matches_full_mix_separator_fallback", report["warnings"])
|
|
self.assertEqual(transcriber.calls, [])
|
|
|
|
def test_near_identical_gain_shifted_full_mix_fallback_is_rejected(self) -> None:
|
|
final, _stem = audio_fixture()
|
|
generator = torch.Generator().manual_seed(91)
|
|
nearly_same = 1.013 * final["waveform"] + 1.0e-5 * torch.randn(
|
|
final["waveform"].shape, generator=generator
|
|
)
|
|
transcriber = FakeTranscriber()
|
|
outputs = self.run_lyrics(
|
|
final,
|
|
{"waveform": nearly_same, "sample_rate": final["sample_rate"]},
|
|
fake_pipeline(transcriber),
|
|
)
|
|
report = json.loads(outputs[0])
|
|
comparison = report["diagnostics"]["stem_full_mix_comparison"]
|
|
self.assertTrue(comparison["near_identical_full_mix"])
|
|
self.assertIn("vocal_stem_matches_full_mix_separator_fallback", report["warnings"])
|
|
self.assertEqual(transcriber.calls, [])
|
|
|
|
def test_valid_vocals_emit_relative_timing_events_and_intervals(self) -> None:
|
|
final, stem = audio_fixture()
|
|
transcriber = FakeTranscriber()
|
|
lyrics = "[Verse 1]\nWe light the night\nAnd hold the sky"
|
|
outputs = self.run_lyrics(final, stem, fake_pipeline(transcriber), lyrics=lyrics)
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["schema"], "diffusiongemma.timed_lyrics_report")
|
|
self.assertEqual(report["version"], 1)
|
|
self.assertEqual(report["analysis_status"], "timing_ready")
|
|
self.assertTrue(report["timing_ready"])
|
|
self.assertTrue(outputs[2])
|
|
self.assertEqual(report["lyrics_sha256"], hashlib.sha256(lyrics.encode("utf-8")).hexdigest())
|
|
self.assertEqual(report["master_audio_sha256"], self.timed._waveform_sha256(final))
|
|
self.assertEqual(report["model_id"], "fake/whisper-tiny")
|
|
self.assertEqual(report["excerpt"], {"start_seconds": 0.0, "duration_seconds": 12.0, "end_seconds": 12.0})
|
|
self.assertEqual(report["vocal_intervals"], [{"start_seconds": 2.0, "end_seconds": 5.0}])
|
|
self.assertEqual(
|
|
report["instrumental_intervals"],
|
|
[
|
|
{"start_seconds": 0.0, "end_seconds": 2.0},
|
|
{"start_seconds": 5.0, "end_seconds": 12.0},
|
|
],
|
|
)
|
|
event = report["events"][0]
|
|
self.assertAlmostEqual(event["start_seconds"], 2.25)
|
|
self.assertAlmostEqual(event["end_seconds"], 3.75)
|
|
self.assertEqual(event["text"], "we light the night")
|
|
self.assertEqual(event["transcript"], "we light the night")
|
|
self.assertEqual(event["authored_lines"], ["We light the night"])
|
|
self.assertEqual(event["confidence"], 1.0)
|
|
self.assertIn("Verified authored-lyric timing", outputs[3])
|
|
|
|
def test_direct_mtb_features_follow_the_loaded_model_dtype(self) -> None:
|
|
class Tokenizer:
|
|
@staticmethod
|
|
def convert_ids_to_tokens(_row):
|
|
return ["<|0.00|>", "hello", "<|1.00|>"]
|
|
|
|
class Processor:
|
|
tokenizer = Tokenizer()
|
|
|
|
@staticmethod
|
|
def __call__(_audio, sampling_rate, return_tensors):
|
|
self.assertEqual(sampling_rate, 16_000)
|
|
self.assertEqual(return_tensors, "pt")
|
|
return {"input_features": torch.ones((1, 80, 300), dtype=torch.float32)}
|
|
|
|
@staticmethod
|
|
def batch_decode(_predicted_ids, skip_special_tokens):
|
|
self.assertTrue(skip_special_tokens)
|
|
return ["hello"]
|
|
|
|
class HalfModel:
|
|
device = torch.device("cpu")
|
|
dtype = torch.float16
|
|
config = SimpleNamespace(max_length=32)
|
|
|
|
@staticmethod
|
|
def generate(input_features, **_kwargs):
|
|
self.assertEqual(input_features.dtype, torch.float16)
|
|
return torch.tensor([[1, 2, 3]], dtype=torch.long)
|
|
|
|
result = self.timed._direct_mtb_transcription(
|
|
{"processor": Processor(), "model": HalfModel()},
|
|
torch.zeros(16_000, dtype=torch.float32),
|
|
16_000,
|
|
"en",
|
|
)
|
|
self.assertEqual(result["text"], "hello")
|
|
self.assertEqual(result["tokens"], ["<|0.00|>", "hello", "<|1.00|>"])
|
|
|
|
def test_instrumental_excerpt_returns_natural_fallback_without_whisper(self) -> None:
|
|
final, stem = audio_fixture(vocal_sections=())
|
|
transcriber = FakeTranscriber()
|
|
outputs = self.run_lyrics(final, stem, fake_pipeline(transcriber))
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertFalse(report["timing_ready"])
|
|
self.assertIn("no_vocal_activity_detected", report["warnings"])
|
|
self.assertEqual(report["vocal_intervals"], [])
|
|
self.assertEqual(
|
|
report["instrumental_intervals"],
|
|
[{"start_seconds": 0.0, "end_seconds": 12.0}],
|
|
)
|
|
self.assertEqual(transcriber.calls, [])
|
|
|
|
def test_weak_authored_alignment_falls_back_instead_of_forcing_words(self) -> None:
|
|
final, stem = audio_fixture()
|
|
transcriber = FakeTranscriber(
|
|
[{"chunks": [{"text": "completely unrelated phrase", "timestamp": [0.0, 1.0]}]}]
|
|
)
|
|
outputs = self.run_lyrics(final, stem, fake_pipeline(transcriber), confidence=0.65)
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertFalse(report["timing_ready"])
|
|
self.assertEqual(report["events"], [])
|
|
self.assertIn("alignment_confidence_below_threshold", report["warnings"])
|
|
|
|
def test_adjacent_whisper_fragments_can_align_to_the_same_authored_line(self) -> None:
|
|
final, stem = audio_fixture()
|
|
transcriber = FakeTranscriber(
|
|
[
|
|
{
|
|
"chunks": [
|
|
{"text": "we light", "timestamp": [0.1, 0.8]},
|
|
{"text": "the night", "timestamp": [0.8, 1.5]},
|
|
]
|
|
}
|
|
]
|
|
)
|
|
outputs = self.run_lyrics(final, stem, fake_pipeline(transcriber))
|
|
report = json.loads(outputs[0])
|
|
self.assertTrue(report["timing_ready"])
|
|
self.assertEqual(len(report["events"]), 2)
|
|
self.assertEqual(
|
|
[event["authored_lines"] for event in report["events"]],
|
|
[["We light the night"], ["We light the night"]],
|
|
)
|
|
|
|
def test_chunk_timestamp_offset_after_fifteen_seconds_is_preserved(self) -> None:
|
|
final, stem = audio_fixture(duration=35.0, vocal_sections=((20.0, 23.0),))
|
|
transcriber = FakeTranscriber(
|
|
[{"chunks": [{"text": "late in the song", "timestamp": [0.4, 1.4]}]}]
|
|
)
|
|
outputs = self.run_lyrics(
|
|
final,
|
|
stem,
|
|
fake_pipeline(transcriber),
|
|
lyrics="[Verse]\nLate in the song",
|
|
)
|
|
report = json.loads(outputs[0])
|
|
self.assertTrue(report["timing_ready"])
|
|
self.assertAlmostEqual(report["events"][0]["start_seconds"], 20.4)
|
|
self.assertAlmostEqual(report["events"][0]["end_seconds"], 21.4)
|
|
self.assertAlmostEqual(transcriber.calls[0]["chunk_offset_seconds"], 20.0)
|
|
self.assertAlmostEqual(transcriber.calls[0]["absolute_song_start_seconds"], 20.0)
|
|
transcription = report["diagnostics"]["transcription"]
|
|
self.assertFalse(transcription["mtb_thirty_second_offset_logic_used"])
|
|
self.assertTrue(transcription["chunk_offsets_added_explicitly"])
|
|
self.assertLessEqual(transcriber.calls[0]["duration"], 15.0)
|
|
|
|
def test_excerpt_start_is_added_for_audio_crop_but_events_remain_excerpt_relative(self) -> None:
|
|
final, stem = audio_fixture(duration=40.0, vocal_sections=((21.0, 24.0),))
|
|
transcriber = FakeTranscriber(
|
|
[{"chunks": [{"text": "inside the window", "timestamp": [0.5, 1.5]}]}]
|
|
)
|
|
outputs = self.run_lyrics(
|
|
final,
|
|
stem,
|
|
fake_pipeline(transcriber),
|
|
lyrics="[Verse]\nInside the window",
|
|
excerpt_start=20.0,
|
|
excerpt_duration=10.0,
|
|
)
|
|
report = json.loads(outputs[0])
|
|
self.assertTrue(report["timing_ready"])
|
|
self.assertEqual(report["vocal_intervals"], [{"start_seconds": 1.0, "end_seconds": 4.0}])
|
|
self.assertAlmostEqual(report["events"][0]["start_seconds"], 1.5)
|
|
self.assertAlmostEqual(transcriber.calls[0]["absolute_song_start_seconds"], 21.0)
|
|
|
|
def test_transcription_exception_falls_back_offloads_model_and_clears_cuda(self) -> None:
|
|
final, stem = audio_fixture()
|
|
transcriber = FakeTranscriber(exception=RuntimeError("synthetic decoder failure"))
|
|
pipeline = fake_pipeline(transcriber, with_model=True)
|
|
with mock.patch.object(self.timed.torch.cuda, "empty_cache") as empty_cache:
|
|
outputs = self.run_lyrics(final, stem, pipeline)
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertFalse(report["timing_ready"])
|
|
self.assertTrue(outputs[2])
|
|
self.assertEqual(report["diagnostics"]["exception_type"], "RuntimeError")
|
|
self.assertEqual(pipeline["model"].moves[-1], "cpu")
|
|
empty_cache.assert_called_once()
|
|
|
|
def test_missing_stem_still_offloads_an_already_loaded_whisper_model(self) -> None:
|
|
final, _stem = audio_fixture()
|
|
pipeline = fake_pipeline(FakeTranscriber(), with_model=True)
|
|
outputs = self.timed.DiffusionGemmaTimedLyricsAnalyzer().analyze(
|
|
final,
|
|
"Lyrics + lip sync",
|
|
"We light the night",
|
|
self.timed._waveform_sha256(final),
|
|
12.0,
|
|
0.0,
|
|
12.0,
|
|
"en",
|
|
0.55,
|
|
None,
|
|
pipeline,
|
|
)
|
|
report = json.loads(outputs[0])
|
|
self.assertIn("vocal_stem_missing", report["warnings"])
|
|
self.assertEqual(pipeline["model"].moves[-1], "cpu")
|
|
|
|
def test_success_is_cached_by_audio_excerpt_lyrics_and_model(self) -> None:
|
|
final, stem = audio_fixture()
|
|
transcriber = FakeTranscriber()
|
|
pipeline = fake_pipeline(transcriber)
|
|
first = self.run_lyrics(final, stem, pipeline)
|
|
second = self.run_lyrics(final, stem, pipeline)
|
|
self.assertTrue(json.loads(first[0])["timing_ready"])
|
|
cached_report = json.loads(second[0])
|
|
self.assertTrue(cached_report["timing_ready"])
|
|
self.assertTrue(cached_report["diagnostics"]["cache_hit"])
|
|
self.assertEqual(len(transcriber.calls), 1)
|
|
self.assertIn("cache", second[1].casefold())
|
|
|
|
def test_nonfinite_or_misaligned_stem_is_never_a_blocker(self) -> None:
|
|
final, stem = audio_fixture()
|
|
bad = {"waveform": stem["waveform"].clone(), "sample_rate": stem["sample_rate"]}
|
|
bad["waveform"][..., 7] = math.nan
|
|
outputs = self.run_lyrics(final, bad, fake_pipeline(FakeTranscriber()))
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertTrue(outputs[2])
|
|
self.assertIn("vocal_stem_invalid_or_misaligned", report["warnings"])
|
|
|
|
shortened = {
|
|
"waveform": stem["waveform"][..., :-1000].clone(),
|
|
"sample_rate": stem["sample_rate"],
|
|
}
|
|
outputs = self.run_lyrics(final, shortened, fake_pipeline(FakeTranscriber()))
|
|
report = json.loads(outputs[0])
|
|
self.assertEqual(report["analysis_status"], "natural_fallback")
|
|
self.assertTrue(outputs[2])
|
|
self.assertIn("vocal_stem_invalid_or_misaligned", report["warnings"])
|
|
|
|
def test_package_registers_analyzer(self) -> None:
|
|
package = load_package()
|
|
self.assertIs(
|
|
package.NODE_CLASS_MAPPINGS["DiffusionGemmaTimedLyricsAnalyzer"],
|
|
sys.modules[f"{package.__name__}.timed_lyrics_nodes"].DiffusionGemmaTimedLyricsAnalyzer,
|
|
)
|
|
self.assertEqual(
|
|
package.NODE_DISPLAY_NAME_MAPPINGS["DiffusionGemmaTimedLyricsAnalyzer"],
|
|
"DiffusionGemma Timed Lyrics Analyzer",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|