123 lines
5 KiB
Python
123 lines
5 KiB
Python
"""Compare H3 audio-latent and lossless-waveform boundaries."""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import wave
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def dbfs(value: float) -> float:
|
|
return 20.0 * math.log10(max(value, 1e-20))
|
|
|
|
|
|
def load_audio_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.float()
|
|
|
|
|
|
def latent_metrics(latent: torch.Tensor) -> dict:
|
|
frames = latent.movedim(-1, 0).flatten(1)
|
|
frame_rms = frames.square().mean(1).sqrt()
|
|
frame_mean = frames.mean(1)
|
|
frame_max = frames.abs().amax(1)
|
|
deltas = frames[1:] - frames[:-1]
|
|
delta_rms = deltas.square().mean(1).sqrt()
|
|
adjacent_cosine = torch.nn.functional.cosine_similarity(frames[:-1], frames[1:], dim=1)
|
|
|
|
block_frames = min(4, frames.shape[0] // 2)
|
|
first = frames[:block_frames].flatten()
|
|
last = frames[-block_frames:].flatten()
|
|
first_last_cosine = torch.nn.functional.cosine_similarity(first, last, dim=0)
|
|
|
|
first_count = min(20, frames.shape[0])
|
|
return {
|
|
"shape": list(latent.shape),
|
|
"dtype": str(latent.dtype),
|
|
"first_20_frame_rms": frame_rms[:first_count].tolist(),
|
|
"first_20_frame_mean": frame_mean[:first_count].tolist(),
|
|
"first_20_frame_max_abs": frame_max[:first_count].tolist(),
|
|
"first_19_delta_rms": delta_rms[: max(0, first_count - 1)].tolist(),
|
|
"first_19_adjacent_cosine": adjacent_cosine[: max(0, first_count - 1)].tolist(),
|
|
"first_4_rms": float(frames[:block_frames].square().mean().sqrt()),
|
|
"frames_4_20_rms": float(frames[block_frames:first_count].square().mean().sqrt()),
|
|
"remaining_rms": float(frames[first_count:].square().mean().sqrt()),
|
|
"first_4_vs_last_4_cosine": float(first_last_cosine),
|
|
"largest_delta_frame": int(delta_rms.argmax().item() + 1),
|
|
"largest_delta_rms": float(delta_rms.max()),
|
|
}
|
|
|
|
|
|
def load_wav(path: Path) -> tuple[np.ndarray, int]:
|
|
with wave.open(str(path), "rb") as source:
|
|
if source.getsampwidth() != 2:
|
|
raise ValueError(f"{path} must be PCM S16")
|
|
channels = source.getnchannels()
|
|
sample_rate = source.getframerate()
|
|
samples = np.frombuffer(source.readframes(source.getnframes()), dtype="<i2")
|
|
return samples.reshape(-1, channels).astype(np.float32) / 32768.0, sample_rate
|
|
|
|
|
|
def waveform_metrics(samples: np.ndarray, sample_rate: int) -> dict:
|
|
first_half_second = samples[: sample_rate // 2]
|
|
mono = first_half_second.mean(1)
|
|
window_samples = sample_rate // 100
|
|
windows = []
|
|
for start in range(0, len(first_half_second), window_samples):
|
|
block = first_half_second[start : start + window_samples]
|
|
if len(block) == 0:
|
|
continue
|
|
windows.append({
|
|
"start_ms": start * 1000.0 / sample_rate,
|
|
"peak_dbfs": dbfs(float(np.max(np.abs(block)))),
|
|
"rms_dbfs": dbfs(float(np.sqrt(np.mean(block * block)))),
|
|
"mean": float(block.mean()),
|
|
})
|
|
|
|
spectrum_samples = min(sample_rate // 10, len(mono))
|
|
windowed = mono[:spectrum_samples] * np.hanning(spectrum_samples)
|
|
magnitudes = np.abs(np.fft.rfft(windowed))
|
|
frequencies = np.fft.rfftfreq(spectrum_samples, 1.0 / sample_rate)
|
|
dominant = np.argsort(magnitudes[1:])[-8:][::-1] + 1
|
|
derivatives = np.max(np.abs(np.diff(first_half_second, axis=0)), axis=1)
|
|
return {
|
|
"sample_rate": sample_rate,
|
|
"samples": len(samples),
|
|
"first_sample": samples[0].tolist(),
|
|
"first_sample_dbfs": [dbfs(float(abs(value))) for value in samples[0]],
|
|
"first_500ms_peak_dbfs": dbfs(float(np.max(np.abs(first_half_second)))),
|
|
"first_500ms_rms_dbfs": dbfs(float(np.sqrt(np.mean(first_half_second**2)))),
|
|
"largest_derivative": float(derivatives.max()),
|
|
"largest_derivative_ms": float((derivatives.argmax() + 1) * 1000.0 / sample_rate),
|
|
"dominant_first_100ms_hz": [float(frequencies[index]) for index in dominant],
|
|
"windows_10ms": windows,
|
|
}
|
|
|
|
|
|
def analyze(latent_path: Path, wav_path: Path) -> dict:
|
|
samples, sample_rate = load_wav(wav_path)
|
|
return {
|
|
"latent_path": str(latent_path),
|
|
"wav_path": str(wav_path),
|
|
"latent": latent_metrics(load_audio_latent(latent_path)),
|
|
"waveform": waveform_metrics(samples, sample_rate),
|
|
}
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--affected-latent", type=Path, required=True)
|
|
parser.add_argument("--affected-wav", type=Path, required=True)
|
|
parser.add_argument("--clean-latent", type=Path, required=True)
|
|
parser.add_argument("--clean-wav", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
print(json.dumps({
|
|
"affected": analyze(args.affected_latent, args.affected_wav),
|
|
"clean": analyze(args.clean_latent, args.clean_wav),
|
|
}, indent=2))
|