tests/test_flow_encoder_pilot.py

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
    )