106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
"""Trace when an H3 audio-boundary artifact emerges during base sampling."""
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig
|
|
from h3_blackwell_runtime.sampler import _decode_audio_latent, sample_video_res_multistep
|
|
from h3_blackwell_runtime.t2v import random_av_latents
|
|
|
|
|
|
def dbfs(value: float) -> float:
|
|
return 20.0 * math.log10(max(value, 1e-20))
|
|
|
|
|
|
def latent_metrics(latent: torch.Tensor) -> dict:
|
|
frames = latent.float().movedim(-1, 0).flatten(1)
|
|
return {
|
|
"first_4_rms": float(frames[:4].square().mean().sqrt()),
|
|
"frames_4_20_rms": float(frames[4:20].square().mean().sqrt()),
|
|
"first_frame_rms": float(frames[0].square().mean().sqrt()),
|
|
"frame_0_to_1_delta_rms": float((frames[1] - frames[0]).square().mean().sqrt()),
|
|
}
|
|
|
|
|
|
def waveform_metrics(waveform: torch.Tensor) -> dict:
|
|
waveform = waveform.float()
|
|
first_100ms = waveform[..., :3200]
|
|
first_500ms = waveform[..., :16000]
|
|
derivative = (first_500ms[..., 1:] - first_500ms[..., :-1]).abs()
|
|
return {
|
|
"first_sample": waveform[..., 0].flatten().tolist(),
|
|
"first_100ms_peak_dbfs": dbfs(float(first_100ms.abs().max())),
|
|
"first_100ms_rms_dbfs": dbfs(float(first_100ms.square().mean().sqrt())),
|
|
"first_500ms_peak_dbfs": dbfs(float(first_500ms.abs().max())),
|
|
"first_500ms_rms_dbfs": dbfs(float(first_500ms.square().mean().sqrt())),
|
|
"largest_derivative": float(derivative.max()),
|
|
}
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--benchmark", type=Path, required=True)
|
|
parser.add_argument("--trace", type=Path, required=True)
|
|
parser.add_argument("--report", type=Path, required=True)
|
|
parser.add_argument("--attention", default="sage2")
|
|
args = parser.parse_args()
|
|
|
|
benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))
|
|
runtime = H3HotRuntime(RuntimeConfig(attention=args.attention))
|
|
video, initial_audio, aligned_frames = random_av_latents(
|
|
benchmark["resolution"][0],
|
|
benchmark["resolution"][1],
|
|
benchmark["frames"],
|
|
benchmark["seed"],
|
|
device=runtime.config.device,
|
|
)
|
|
text = runtime.refiner(runtime.conditioner(benchmark["prompt"]))
|
|
trace = []
|
|
video, final_audio = sample_video_res_multistep(
|
|
runtime.model,
|
|
runtime.packer,
|
|
text,
|
|
video,
|
|
initial_audio,
|
|
steps=benchmark["steps"],
|
|
seed=benchmark["seed"],
|
|
return_audio=True,
|
|
audio_step_trace=trace,
|
|
)
|
|
|
|
report = {
|
|
"benchmark": str(args.benchmark),
|
|
"attention": args.attention,
|
|
"seed": benchmark["seed"],
|
|
"frames": aligned_frames,
|
|
"steps": [],
|
|
}
|
|
with torch.inference_mode():
|
|
for entry in trace:
|
|
denoised = _decode_audio_latent(entry["audio_denoised"]).to(
|
|
"cuda", dtype=next(runtime.audio_vae.parameters()).dtype,
|
|
)
|
|
waveform = runtime.audio_vae.decode(denoised).cpu()[0]
|
|
report["steps"].append({
|
|
"step": entry["step"],
|
|
"video_sigma": entry["video_sigma"],
|
|
"audio_sigma": entry["audio_sigma"],
|
|
"latent": latent_metrics(denoised.cpu()),
|
|
"denoised_waveform": waveform_metrics(waveform),
|
|
})
|
|
|
|
args.trace.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({
|
|
"initial_audio": initial_audio.detach().cpu(),
|
|
"final_audio": final_audio.detach().cpu(),
|
|
"steps": trace,
|
|
"prompt": benchmark["prompt"],
|
|
"seed": benchmark["seed"],
|
|
}, args.trace)
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
serialized = json.dumps(report, indent=2)
|
|
args.report.write_text(serialized + "\n", encoding="utf-8")
|
|
print(serialized)
|