182 lines
7.1 KiB
Python
182 lines
7.1 KiB
Python
"""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)
|