h3-blackwell-runtime/tools/probe_audio_vae_boundary.py

"""Decode controlled H3 audio-latent boundary variants for diagnosis."""

import argparse
import json
import math
import subprocess
from pathlib import Path

import numpy as np
import torch

from h3_blackwell_runtime.audio_vae_decoder import MiniMaxH3AudioVAE


SAMPLE_RATE = 32000
SAMPLES_PER_LATENT = 800


def load_latent(path: Path) -> torch.Tensor:
    state = torch.load(path, map_location="cpu", weights_only=False)
    latent = state.get("audio_latent") if isinstance(state, dict) else state
    if latent is None or latent.ndim != 4:
        raise ValueError(f"{path} does not contain a [B,C,S,T] audio latent")
    return latent


def dbfs(value: float) -> float:
    return 20.0 * math.log10(max(value, 1e-20))


def waveform_metrics(waveform: torch.Tensor) -> dict:
    samples = waveform.float().numpy().T
    first_100ms = samples[: SAMPLE_RATE // 10]
    first_500ms = samples[: SAMPLE_RATE // 2]
    derivatives = np.max(np.abs(np.diff(first_500ms, axis=0)), axis=1)
    return {
        "samples": len(samples),
        "first_sample": samples[0].tolist(),
        "first_sample_dbfs": [dbfs(float(abs(value))) for value in samples[0]],
        "first_100ms_peak_dbfs": dbfs(float(np.max(np.abs(first_100ms)))),
        "first_100ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_100ms**2)))),
        "first_500ms_peak_dbfs": dbfs(float(np.max(np.abs(first_500ms)))),
        "first_500ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_500ms**2)))),
        "largest_derivative": float(derivatives.max()),
        "largest_derivative_ms": float((derivatives.argmax() + 1) * 1000.0 / SAMPLE_RATE),
    }


def comparison(reference: torch.Tensor, candidate: torch.Tensor) -> dict:
    count = min(reference.shape[-1], candidate.shape[-1])
    reference = reference[..., :count].float()
    candidate = candidate[..., :count].float()

    def region_metrics(samples: int) -> dict:
        ref = reference[..., :samples]
        test = candidate[..., :samples]
        error = test - ref
        signal_power = ref.square().mean(dim=-1)
        noise_power = error.square().mean(dim=-1)
        psnr = 10.0 * torch.log10(signal_power.clamp_min(1e-30) / noise_power.clamp_min(1e-30))
        return {
            "rmse": float(error.square().mean().sqrt()),
            "max_abs": float(error.abs().max()),
            "psnr_db_by_channel": psnr.flatten().tolist(),
        }

    return {
        "first_100ms": region_metrics(SAMPLE_RATE // 10),
        "first_500ms": region_metrics(SAMPLE_RATE // 2),
        "full": region_metrics(count),
    }


def write_waveform(path: Path, waveform: torch.Tensor) -> None:
    raw = path.with_suffix(".f32le")
    waveform.transpose(0, 1).contiguous().numpy().tofile(raw)
    subprocess.run([
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
        "-f", "f32le", "-ar", str(SAMPLE_RATE), "-ac", "2", "-i", str(raw),
        "-c:a", "pcm_f32le", str(path),
    ], check=True)
    raw.unlink()


parser = argparse.ArgumentParser()
parser.add_argument("--affected-latent", type=Path, required=True)
parser.add_argument("--clean-latent", type=Path, required=True)
parser.add_argument("--vae", type=Path, default=Path("/vae/minimax_h3_audio_vae_fp32.safetensors"))
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--report", type=Path)
args = parser.parse_args()

affected = load_latent(args.affected_latent)
clean = load_latent(args.clean_latent)
if affected.shape != clean.shape:
    raise ValueError(f"latent shapes differ: {tuple(affected.shape)} != {tuple(clean.shape)}")

boundary_frames = 4
variants = {
    "affected-original": (affected, 0),
    "clean-original": (clean, 0),
    "zero-normalized-latent": (torch.zeros_like(affected), 0),
    "affected-repeat-frame0": (affected[..., :1].expand_as(affected).clone(), 0),
    "affected-repeat-frame4": (affected[..., 4:5].expand_as(affected).clone(), 0),
}

silent_carrier = affected[..., 4:5].expand_as(affected).clone()
carrier_start = silent_carrier.clone()
carrier_start[..., :boundary_frames] = affected[..., :boundary_frames]
variants["carrier-affected-first4-at-start"] = (carrier_start, 0)

interior_frame = 40
carrier_interior = silent_carrier.clone()
carrier_interior[..., interior_frame : interior_frame + boundary_frames] = affected[..., :boundary_frames]
variants["carrier-affected-first4-at-frame40"] = (carrier_interior, 0)

replaced_with_frame4 = affected.clone()
replaced_with_frame4[..., :boundary_frames] = affected[..., 4:5]
variants["affected-first4-repeat-frame4"] = (replaced_with_frame4, 0)

affected_with_clean = affected.clone()
affected_with_clean[..., :boundary_frames] = clean[..., :boundary_frames]
variants["affected-first4-from-clean"] = (affected_with_clean, 0)

clean_with_affected = clean.clone()
clean_with_affected[..., :boundary_frames] = affected[..., :boundary_frames]
variants["clean-first4-from-affected"] = (clean_with_affected, 0)

prefix_repeat = affected[..., :1].expand(*affected.shape[:-1], boundary_frames)
variants["affected-prefix-repeat-frame0"] = (
    torch.cat((prefix_repeat, affected), dim=-1),
    boundary_frames * SAMPLES_PER_LATENT,
)
variants["affected-prefix-own-first4"] = (
    torch.cat((affected[..., :boundary_frames], affected), dim=-1),
    boundary_frames * SAMPLES_PER_LATENT,
)

vae = MiniMaxH3AudioVAE.from_safetensors(args.vae, device="cuda").eval()
args.output_dir.mkdir(parents=True, exist_ok=True)
decoded = {}
report = {"boundary_frames": boundary_frames, "variants": {}}
with torch.inference_mode():
    for name, (latent, crop_start) in variants.items():
        waveform = vae.decode(latent.to("cuda", dtype=next(vae.parameters()).dtype)).cpu()[0]
        waveform = waveform[:, crop_start : crop_start + affected.shape[-1] * SAMPLES_PER_LATENT]
        decoded[name] = waveform
        output = args.output_dir / f"{name}.wav"
        write_waveform(output, waveform)
        report["variants"][name] = {
            "output": str(output),
            "crop_start_samples": crop_start,
            "metrics": waveform_metrics(waveform),
        }

affected_reference = decoded["affected-original"]
for name, waveform in decoded.items():
    if name != "affected-original":
        report["variants"][name]["difference_from_affected_original"] = comparison(
            affected_reference, waveform,
        )

interior_start = interior_frame * SAMPLES_PER_LATENT
segment_samples = boundary_frames * SAMPLES_PER_LATENT
report["interior_placement"] = {
    "frame": interior_frame,
    "start_seconds": interior_start / SAMPLE_RATE,
    "affected_onset_vs_carrier_interior_event": comparison(
        affected_reference[..., :segment_samples],
        decoded["carrier-affected-first4-at-frame40"][..., interior_start : interior_start + segment_samples],
    ),
    "carrier_start_event_vs_carrier_interior_event": comparison(
        decoded["carrier-affected-first4-at-start"][..., :segment_samples],
        decoded["carrier-affected-first4-at-frame40"][..., interior_start : interior_start + segment_samples],
    ),
}

serialized = json.dumps(report, indent=2)
if args.report is not None:
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(serialized + "\n", encoding="utf-8")
print(serialized)