| from __future__ import annotations |
|
|
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from music3lab.codec.flow_encoder import ( |
| ContinuousFlowEncoder, |
| audio_metrics, |
| latent_normalized_mse, |
| learning_rate, |
| load_flow_encoder_config, |
| prior_mean_latents, |
| ) |
| from music3lab.codec.runner import ( |
| FileRecord, |
| TeacherDatasetManifest, |
| TeacherSplitRecord, |
| ) |
| from music3lab.vocoder import FlowVocoderLatents |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| CONFIG = ROOT / "configs" / "flow-encoder-v1.yaml" |
|
|
|
|
| class TinyVocoder(nn.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.gain = nn.Parameter(torch.ones(2), requires_grad=False) |
|
|
| def forward(self, latents: FlowVocoderLatents) -> torch.Tensor: |
| value = F.interpolate( |
| latents.tensor[:, :2].float(), |
| size=44032, |
| mode="linear", |
| align_corners=False, |
| ) |
| return value * self.gain.view(1, 2, 1) |
|
|
|
|
| def test_frozen_config_has_exact_geometry_and_disjoint_64_16_16_seeds() -> None: |
| config = load_flow_encoder_config(CONFIG).config |
| seeds = config.teacher.split_seeds() |
| assert {key: len(value) for key, value in seeds.items()} == { |
| "train": 64, |
| "validation": 16, |
| "heldout": 16, |
| } |
| assert len({item for values in seeds.values() for item in values}) == 96 |
| assert config.capability.endswith("not_native_rvq") |
| assert config.model.input_samples == 44032 |
| assert config.model.latent_frames == 86 |
|
|
|
|
| def test_encoder_exact_geometry_is_deterministic_and_has_gradients() -> None: |
| loaded = load_flow_encoder_config(CONFIG).config |
| torch.manual_seed(9) |
| encoder = ContinuousFlowEncoder( |
| loaded.model, latent_mean=loaded.loss.latent_mean |
| ) |
| audio = torch.randn(2, 2, 44032) |
| first = encoder(audio) |
| second = encoder(audio) |
| assert first.shape == (2, 128, 86) |
| assert torch.equal(first, second) |
| target = torch.randn_like(first) |
| latent_normalized_mse(first, target, loaded.loss).mean().backward() |
| assert any( |
| parameter.grad is not None and parameter.grad.abs().sum() > 0 |
| for parameter in encoder.parameters() |
| ) |
|
|
|
|
| def test_prior_mean_baseline_and_schedule_are_frozen() -> None: |
| config = load_flow_encoder_config(CONFIG).config |
| prior = prior_mean_latents(3, config.model, config.loss) |
| assert prior.shape == (3, 128, 86) |
| assert float(prior.mean()) == pytest.approx(config.loss.latent_mean) |
| assert learning_rate(0, config.training) == pytest.approx( |
| config.training.maximum_learning_rate |
| ) |
| assert learning_rate(config.training.steps, config.training) == pytest.approx( |
| config.training.minimum_learning_rate |
| ) |
|
|
|
|
| def test_frozen_decoder_has_no_weight_grad_while_encoder_input_does() -> None: |
| decoder = TinyVocoder() |
| latents = torch.randn(2, 128, 86, requires_grad=True) |
| audio = decoder(FlowVocoderLatents(latents)) |
| audio.square().mean().backward() |
| assert latents.grad is not None and latents.grad.abs().sum() > 0 |
| assert all(parameter.grad is None for parameter in decoder.parameters()) |
|
|
|
|
| def test_audio_metrics_order_exact_before_distorted() -> None: |
| target = torch.randn(2, 2, 44032) |
| noise = 0.1 * torch.randn(2, 2, 44032) |
| exact = audio_metrics(target, target) |
| distorted = audio_metrics(target + noise, target) |
| assert exact.mae == 0 |
| assert exact.correlation > distorted.correlation |
| assert exact.unscaled_snr_db > distorted.unscaled_snr_db |
|
|
|
|
| def test_cli_help_is_cpu_safe() -> None: |
| result = subprocess.run( |
| [ |
| sys.executable, |
| "-B", |
| str(ROOT / "scripts" / "run_flow_encoder_pilot.py"), |
| "--help", |
| ], |
| check=True, |
| text=True, |
| capture_output=True, |
| env={ |
| "PATH": __import__("os").environ["PATH"], |
| "PYTHONPATH": str(ROOT / "src"), |
| "CUDA_VISIBLE_DEVICES": "-1", |
| "PYTHONDONTWRITEBYTECODE": "1", |
| }, |
| ) |
| normalized = " ".join(result.stdout.split()) |
| assert "does not produce native RVQ tokens" in normalized |
|
|
|
|
| def test_teacher_manifest_canonicalizes_nested_split_models() -> None: |
| config = load_flow_encoder_config(CONFIG) |
| counts = {"train": 64, "validation": 16, "heldout": 16} |
| starts = {"train": 1000, "validation": 2000, "heldout": 3000} |
| splits = { |
| name: TeacherSplitRecord( |
| count=count, |
| seeds=tuple(range(starts[name], starts[name] + count)), |
| file=FileRecord( |
| path=f"{name}.safetensors", |
| sha256="1" * 64, |
| size=1, |
| ), |
| audio_shape=(count, 2, 44032), |
| latent_shape=(count, 128, 86), |
| audio_content_sha256="2" * 64, |
| latent_content_sha256="3" * 64, |
| ) |
| for name, count in counts.items() |
| } |
| manifest = TeacherDatasetManifest.create( |
| schema_version="music3lab.flow-encoder-teachers.v1", |
| capability="continuous_flow_vocoder_latent_teacher_pairs_not_native_rvq", |
| config_file_sha256=config.file_sha256, |
| config_semantic_digest=config.semantic_digest, |
| model_revision=config.config.model_revision, |
| base_id=config.config.expected_base_id, |
| base_manifest_file_sha256="4" * 64, |
| diffusers_revision=config.config.diffusers_revision, |
| prompt_sha256="5" * 64, |
| lyrics_sha256="6" * 64, |
| persistent_pipeline_load_count=1, |
| replay_exact_count=96, |
| producer_project_git_commit="7" * 40, |
| publication_project_git_commit="8" * 40, |
| recovered_from_complete_quarantine=True, |
| generation_seconds=None, |
| peak_cuda_allocated_bytes=None, |
| splits=splits, |
| ) |
| assert ( |
| TeacherDatasetManifest.model_validate_json( |
| __import__( |
| "music3lab.codec.flow_encoder", |
| fromlist=["canonical_json_bytes"], |
| ).canonical_json_bytes(manifest) |
| ) |
| == manifest |
| ) |
|
|
