Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
486 lines
27 KiB
Python
486 lines
27 KiB
Python
#!/usr/bin/env python3
|
|
"""Capture AuK reference boundaries for ordered native parity testing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import random
|
|
import sys
|
|
import time
|
|
import statistics
|
|
|
|
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
|
|
|
|
import numpy as np
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--reference", type=Path, default=Path("reference/AuK"))
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--qwen-dir", type=Path, required=True)
|
|
parser.add_argument("--instruction", required=True)
|
|
parser.add_argument("--audio", type=Path)
|
|
parser.add_argument("--duration", type=float, required=True)
|
|
parser.add_argument("--steps", type=int, default=32)
|
|
parser.add_argument("--cfg", type=float, default=2.0)
|
|
parser.add_argument("--sway", type=float, default=-1.0)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16")
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--log", action="store_true", required=True)
|
|
parser.add_argument("--capture-layers", action="store_true")
|
|
parser.add_argument("--qkv-fixture", type=Path)
|
|
parser.add_argument("--audio-replay-fixture", type=Path)
|
|
parser.add_argument("--audio-features", type=Path)
|
|
parser.add_argument("--full-precision-reduction", action="store_true")
|
|
parser.add_argument("--decode-fixture", type=Path)
|
|
parser.add_argument("--flow-fixture", type=Path)
|
|
parser.add_argument("--flow-cfg", action="store_true")
|
|
parser.add_argument("--flow-sample", action="store_true")
|
|
parser.add_argument("--benchmark-runs", type=int, default=0)
|
|
args = parser.parse_args()
|
|
if bool(args.audio_replay_fixture) != bool(args.audio_features):
|
|
parser.error("--audio-replay-fixture and --audio-features must be supplied together")
|
|
if args.decode_fixture and args.flow_fixture:
|
|
parser.error("choose only one component fixture")
|
|
if args.flow_cfg and not args.flow_fixture:
|
|
parser.error("--flow-cfg requires --flow-fixture")
|
|
if args.flow_sample and not args.flow_fixture:
|
|
parser.error("--flow-sample requires --flow-fixture")
|
|
if args.benchmark_runs < 0 or (args.benchmark_runs and args.flow_fixture):
|
|
parser.error("--benchmark-runs requires a nonnegative count and VAE-only or full inference")
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
handlers=[logging.FileHandler(args.output_dir / "reference.log"), logging.StreamHandler()],
|
|
force=True,
|
|
)
|
|
sys.path.insert(0, str(args.reference.resolve() / "src"))
|
|
from auk.infer.infer_auk import AukInfer, save_audio
|
|
|
|
torch.set_num_threads(8)
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
torch.backends.cudnn.benchmark = False
|
|
torch.backends.cudnn.deterministic = True
|
|
if args.full_precision_reduction:
|
|
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
|
|
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
|
|
if args.flow_fixture:
|
|
from auk.model.flux2_edit import Flux2Edit
|
|
from omegaconf import OmegaConf
|
|
manifest = json.loads((args.flow_fixture / "manifest.json").read_text())
|
|
if manifest["dtype"] != "fp32":
|
|
raise ValueError("flow embedding fixture currently targets FP32 only")
|
|
config = OmegaConf.load(args.model_dir / "config.yaml")
|
|
arch = OmegaConf.to_container(config.model.arch, resolve=True)
|
|
arch["attn_backend"] = "torch"
|
|
model = Flux2Edit(**arch, latent_dim=64)
|
|
checkpoint = load_file(args.model_dir / "auk_base.safetensors")
|
|
model.load_state_dict({key.removeprefix("transformer."): value
|
|
for key, value in checkpoint.items() if key.startswith("transformer.")}, strict=True)
|
|
del checkpoint
|
|
model = model.cuda().eval()
|
|
fixture = load_file(args.flow_fixture / "reference.safetensors")
|
|
reference_latents = fixture["fm.first.ref"].cuda()
|
|
reference_mask = fixture["fm.first.ref_mask"].cuda() if reference_latents.shape[1] else None
|
|
if not fixture["fm.first.c_mask"].all():
|
|
raise ValueError("flow fixture currently requires all-valid text tokens")
|
|
if args.flow_sample:
|
|
from torchdiffeq import odeint
|
|
steps, strength = manifest["steps"], manifest["cfg"]
|
|
times = torch.linspace(0, 1, steps + 1, device="cuda", dtype=torch.float32)
|
|
times = times + manifest.get("sway", -1.0) * (torch.cos(torch.pi / 2 * times) - 1 + times)
|
|
text = fixture["fm.first.text"].cuda()
|
|
mask = fixture["fm.first.c_mask"].cuda()
|
|
states = {"times": times.cpu()}
|
|
with torch.inference_mode():
|
|
def velocity(time, latent):
|
|
prediction = model(x=latent, text=text, time=time, c_mask=mask, cfg_infer=True, cache=True,
|
|
ref=reference_latents, ref_mask=reference_mask)
|
|
conditional, unconditional = prediction.chunk(2, dim=0)
|
|
return conditional + (conditional - unconditional) * strength
|
|
trajectory = odeint(velocity, fixture["fm.first.x"].cuda(), times, method="euler")
|
|
states["trajectory"] = trajectory.cpu().contiguous()
|
|
states["fm.output"] = trajectory[-1].cpu().contiguous().clone()
|
|
expected = fixture["vae.normalized_input"] if reference_latents.shape[1] else fixture["fm.output"]
|
|
error = (states["fm.output"] - expected).abs().max().item()
|
|
logging.info("Flow-only replay versus full Python output max_error=%g", error)
|
|
torch.testing.assert_close(states["fm.output"], expected, rtol=0, atol=0)
|
|
save_file(states, args.output_dir / "flow-trajectory.safetensors")
|
|
return
|
|
stages = {}
|
|
with torch.inference_mode():
|
|
stages["text"] = model.project_text(fixture["fm.first.text"].cuda()).cpu().contiguous()
|
|
stages["audio"] = model.audio_embed(fixture["fm.first.x"].cuda()).cpu().contiguous()
|
|
if reference_latents.shape[1]:
|
|
stages["reference"] = model.audio_embed._embed(reference_latents, reference_mask).cpu().contiguous()
|
|
stages["reference.unconditional"] = model.audio_embed._embed(
|
|
torch.zeros_like(reference_latents), reference_mask).cpu().contiguous()
|
|
for index, step in enumerate((0.0, 0.25, 0.75)):
|
|
time_embedding = model.time_embed(torch.tensor([step], device="cuda"))
|
|
stages[f"time.{index}"] = time_embedding.cpu().contiguous()
|
|
audio_out, audio_mask, prompt_len = model._embed_audio(
|
|
fixture["fm.first.x"].cuda(), reference_latents, False, None, reference_mask)
|
|
text_out = stages["text"].cuda()
|
|
text_mask = fixture["fm.first.c_mask"].cuda()
|
|
if args.flow_cfg:
|
|
time_embedding = torch.cat((time_embedding, time_embedding), dim=0)
|
|
unconditional_audio, _, _ = model._embed_audio(
|
|
fixture["fm.first.x"].cuda(), reference_latents, True, None, reference_mask)
|
|
audio_out = torch.cat((audio_out, unconditional_audio), dim=0)
|
|
if audio_mask is not None:
|
|
audio_mask = torch.cat((audio_mask, audio_mask), dim=0)
|
|
text_mask = torch.cat((text_mask, text_mask), dim=0)
|
|
text_out = torch.cat((text_out, torch.zeros_like(text_out)), dim=0)
|
|
handles = []
|
|
first = model.transformer_blocks[0]
|
|
for branch in ("x", "c"):
|
|
suffix = "" if branch == "x" else "_c"
|
|
for label, module in (
|
|
("mod", getattr(first, "attn_norm_" + branch)),
|
|
("qkv", getattr(first.attn, "to_qkv" + suffix)),
|
|
("attn_output", first.attn.to_out[0] if branch == "x" else first.attn.to_out_c),
|
|
("q_norm", first.attn.q_norm if branch == "x" else first.attn.c_q_norm),
|
|
("k_norm", first.attn.k_norm if branch == "x" else first.attn.c_k_norm),
|
|
("ff_output", getattr(first, "ff_" + branch)),
|
|
):
|
|
name = f"first.{branch}.{label}.{index}"
|
|
def capture(module, inputs, output, name=name, label=label):
|
|
value = output[0] if isinstance(output, tuple) else output
|
|
if label in ("q_norm", "k_norm"):
|
|
value = value.transpose(1, 2)
|
|
stages[name] = value.detach().cpu().contiguous()
|
|
handles.append(module.register_forward_hook(capture))
|
|
name = f"first.{branch}.ff_input.{index}"
|
|
def capture_input(module, inputs, name=name):
|
|
stages[name] = inputs[0].detach().cpu().contiguous()
|
|
handles.append(getattr(first, "ff_" + branch).register_forward_pre_hook(capture_input))
|
|
name = f"first.{branch}.attention.{index}"
|
|
def capture_attention(module, inputs, name=name):
|
|
stages[name] = inputs[0].detach().cpu().contiguous()
|
|
projection = first.attn.to_out[0] if branch == "x" else first.attn.to_out_c
|
|
handles.append(projection.register_forward_pre_hook(capture_attention))
|
|
for layer, block in enumerate(model.transformer_blocks):
|
|
text_out, audio_out = block(
|
|
audio_out, text_out, time_embedding,
|
|
mask=audio_mask, c_mask=text_mask,
|
|
rope=model.rotary_embed.forward_from_seq_len(audio_out.shape[1]),
|
|
c_rope=model.rotary_embed.forward_from_seq_len(stages["text"].shape[1]))
|
|
if layer == 0:
|
|
stages[f"block.audio.{index}"] = audio_out.cpu().contiguous()
|
|
stages[f"block.text.{index}"] = text_out.cpu().contiguous()
|
|
for handle in handles:
|
|
handle.remove()
|
|
stages[f"double.audio.{index}"] = audio_out.cpu().contiguous()
|
|
stages[f"double.text.{index}"] = text_out.cpu().contiguous()
|
|
hidden = torch.cat((text_out, audio_out), dim=1)
|
|
rope = model.rotary_embed.forward_from_seq_len(hidden.shape[1])
|
|
single_mask = torch.cat((text_mask, audio_mask), dim=1) if audio_mask is not None else None
|
|
for block in model.single_transformer_blocks:
|
|
hidden = block(hidden, time_embedding, mask=single_mask, rope=rope)
|
|
stages[f"single.{index}"] = hidden.cpu().contiguous()
|
|
projected = model.proj_out(model.norm_out(hidden[:, text_out.shape[1] + prompt_len:], time_embedding))
|
|
direct = model(x=fixture["fm.first.x"].cuda(), text=fixture["fm.first.text"].cuda(),
|
|
time=torch.tensor([step], device="cuda"),
|
|
c_mask=fixture["fm.first.c_mask"].cuda(), cfg_infer=args.flow_cfg,
|
|
ref=reference_latents, ref_mask=reference_mask)
|
|
torch.testing.assert_close(projected, direct, rtol=0, atol=0)
|
|
stages[f"velocity.{index}"] = direct.cpu().contiguous()
|
|
if args.flow_cfg:
|
|
conditional, unconditional = direct.chunk(2, dim=0)
|
|
stages[f"conditional.{index}"] = conditional.cpu().contiguous().clone()
|
|
stages[f"unconditional.{index}"] = unconditional.cpu().contiguous().clone()
|
|
stages[f"guided.{index}"] = (conditional + (conditional - unconditional) * 2.0).cpu().contiguous()
|
|
save_file(stages, args.output_dir / ("flow-cfg.safetensors" if args.flow_cfg else "flow-embeddings.safetensors"))
|
|
logging.info("Saved FP32 flow boundaries and velocity at three timesteps (CFG=%s); composition matches original forward exactly", args.flow_cfg)
|
|
return
|
|
|
|
if args.decode_fixture:
|
|
from auk.model.vae import BigVGANFlowVAEConfig, load_vae_model
|
|
from omegaconf import OmegaConf
|
|
config = OmegaConf.load(args.model_dir / "config.yaml")
|
|
vae_config = BigVGANFlowVAEConfig.from_dict(
|
|
OmegaConf.to_container(config.model.vae.model_init_kwargs, resolve=True))
|
|
vae = load_vae_model("BigVGANFlowVAE", vae_config,
|
|
str(args.model_dir / "vae.safetensors"), map_location="cpu").cuda().eval()
|
|
fixture = load_file(args.decode_fixture / "reference.safetensors")
|
|
stages = {}
|
|
handles = []
|
|
def save_stage(name):
|
|
def hook(_module, _inputs, output):
|
|
stages[name] = output.detach().cpu().contiguous().clone()
|
|
return hook
|
|
names = ("conv_pre", "ups.0.0", "resblocks.0.activations.0", "resblocks.0.convs1.0")
|
|
for name, module in vae.named_modules():
|
|
if name in names:
|
|
handles.append(module.register_forward_hook(save_stage(name)))
|
|
with torch.inference_mode():
|
|
latents = vae.denormalize(fixture["fm.output"].cuda()).transpose(1, 2).contiguous()
|
|
stages["input"] = latents.cpu()
|
|
stages["waveform"] = vae.inference_from_latents(latents).cpu()
|
|
for handle in handles:
|
|
handle.remove()
|
|
logging.info("Decoder rerun max waveform error: %.9g",
|
|
(stages["waveform"].flatten() - fixture["waveform"].flatten()).abs().max())
|
|
save_file(stages, args.output_dir / "vae-stages.safetensors")
|
|
if args.benchmark_runs:
|
|
milliseconds = []
|
|
with torch.inference_mode():
|
|
# Match native timing: CPU normalized latents through CPU waveform;
|
|
# exclude weight loading, graph construction and WAV serialization.
|
|
for run in range(args.benchmark_runs + 2):
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
decoder_input = vae.denormalize(fixture["fm.output"].cuda()).transpose(1, 2)
|
|
waveform = vae.inference_from_latents(decoder_input).cpu()
|
|
torch.cuda.synchronize()
|
|
elapsed = (time.perf_counter() - start) * 1000
|
|
if not torch.equal(waveform, stages["waveform"]):
|
|
raise RuntimeError("benchmark VAE output changed")
|
|
if run >= 2:
|
|
milliseconds.append(elapsed)
|
|
logging.info("VAE decode_ms=%.6f", elapsed)
|
|
result = {"runs": args.benchmark_runs, "milliseconds": milliseconds,
|
|
"median_ms": statistics.median(milliseconds), "weights": "original F32",
|
|
"scope": "CPU normalized latents to CPU waveform; no capture hooks"}
|
|
(args.output_dir / "vae-performance.json").write_text(json.dumps(result, indent=2) + "\n")
|
|
logging.info("VAE median_ms=%.6f runs=%d", result["median_ms"], args.benchmark_runs)
|
|
return
|
|
|
|
engine = AukInfer(
|
|
str(args.model_dir / "config.yaml"),
|
|
str(args.model_dir / "auk_base.safetensors"),
|
|
device="cuda",
|
|
dtype=args.dtype,
|
|
qwen_path=str(args.qwen_dir),
|
|
cpu_offload=False,
|
|
)
|
|
if args.audio_replay_fixture:
|
|
if args.dtype != "fp32":
|
|
raise ValueError("audio feature replay requires FP32")
|
|
fixture = load_file(args.audio_replay_fixture / "reference.safetensors")
|
|
lengths = fixture["qwen.audio.feature_lens"].cuda()
|
|
if lengths.numel() != 1:
|
|
raise ValueError("audio feature replay requires one reference")
|
|
frames = int(lengths.item())
|
|
original = fixture["condition.input_features"][0, :, :frames].cuda()
|
|
native = load_file(args.audio_features)["features"].cuda()
|
|
if native.shape != original.shape:
|
|
raise ValueError("replayed audio feature shape differs")
|
|
expected = fixture["qwen.audio_output"].cuda()
|
|
with torch.inference_mode():
|
|
outputs = {}
|
|
for name, features in (("original", original), ("native_features", native)):
|
|
output = engine.model.text_encoder.audio_tower(
|
|
features, feature_lens=lengths).last_hidden_state
|
|
x, y = output.double().flatten(), expected.double().flatten()
|
|
cosine = torch.dot(x, y) / (x.norm() * y.norm())
|
|
rmse = (x - y).square().mean().sqrt()
|
|
logging.info("Python audio replay %s: cosine=%.12g RMSE=%.12g",
|
|
name, cosine.item(), rmse.item())
|
|
outputs[name] = output.cpu().contiguous()
|
|
save_file(outputs, args.output_dir / "audio-replay.safetensors")
|
|
return
|
|
captured: dict[str, torch.Tensor] = {}
|
|
|
|
def capture_inputs(prefix):
|
|
def hook(_module, _args, kwargs):
|
|
for name, value in kwargs.items():
|
|
key = f"{prefix}.{name}"
|
|
if isinstance(value, torch.Tensor) and key not in captured:
|
|
captured[key] = value.detach().cpu().contiguous().clone()
|
|
return hook
|
|
|
|
handles = [
|
|
engine.model.text_encoder.register_forward_pre_hook(capture_inputs("condition"), with_kwargs=True),
|
|
engine.model.text_encoder.model.register_forward_pre_hook(capture_inputs("qwen"), with_kwargs=True),
|
|
engine.model.transformer.register_forward_pre_hook(capture_inputs("fm.first"), with_kwargs=True),
|
|
]
|
|
if args.qkv_fixture:
|
|
if not args.capture_layers or args.benchmark_runs:
|
|
raise ValueError("QKV replay requires layer capture and cannot be benchmarked")
|
|
qkv_fixture = load_file(args.qkv_fixture)
|
|
def replay_projection(name):
|
|
def hook(_module, _inputs, output):
|
|
replacement = qkv_fixture[name].to(device=output.device, dtype=output.dtype)
|
|
if replacement.shape != output.shape:
|
|
raise ValueError(f"QKV replay shape mismatch: {name}")
|
|
return replacement
|
|
return hook
|
|
for name in ("q_proj", "k_proj", "v_proj"):
|
|
handles.append(getattr(engine.model.text_encoder.model.layers[0].self_attn, name)
|
|
.register_forward_hook(replay_projection(name)))
|
|
if args.capture_layers:
|
|
def capture_layer(index):
|
|
def hook(_module, _inputs, output):
|
|
captured[f"qwen.layer.{index}"] = output.detach().cpu().contiguous().clone()
|
|
return hook
|
|
for index, layer in enumerate(engine.model.text_encoder.model.layers):
|
|
handles.append(layer.register_forward_hook(capture_layer(index)))
|
|
for name in ("q_proj", "k_proj", "v_proj"):
|
|
handles.append(getattr(engine.model.text_encoder.model.layers[0].self_attn, name)
|
|
.register_forward_hook(capture_layer(name)))
|
|
def capture_attention_input(_module, inputs):
|
|
captured["qwen.attention_output"] = inputs[0].detach().cpu().contiguous().clone()
|
|
handles.append(engine.model.text_encoder.model.layers[0].self_attn.o_proj
|
|
.register_forward_pre_hook(capture_attention_input))
|
|
original_sample = engine.model.sample
|
|
|
|
def capture_sample(*sample_args, **sample_kwargs):
|
|
for name, value in sample_kwargs.items():
|
|
if isinstance(value, torch.Tensor):
|
|
captured[f"fm.sample.{name}"] = value.detach().cpu().contiguous().clone()
|
|
output, trajectory = original_sample(*sample_args, **sample_kwargs)
|
|
captured["fm.output"] = output.detach().cpu().contiguous()
|
|
return output, trajectory
|
|
|
|
engine.model.sample = capture_sample
|
|
original_denormalize = engine.vae_model.denormalize
|
|
def capture_denormalize(latents):
|
|
captured["vae.normalized_input"] = latents.detach().cpu().contiguous().clone()
|
|
return original_denormalize(latents)
|
|
engine.vae_model.denormalize = capture_denormalize
|
|
original_encoding = engine.vae_model.encoding_and_normalization
|
|
if args.audio:
|
|
def capture_encoder(_module, inputs, output):
|
|
captured["vae.encoder_input"] = inputs[0].detach().cpu().contiguous().clone()
|
|
captured["vae.encoder_stats"] = output.detach().cpu().contiguous().clone()
|
|
handles.append(engine.vae_model.audio_encoder.register_forward_hook(capture_encoder))
|
|
def capture_encoding(*encoding_args, **encoding_kwargs):
|
|
original_randn_like = torch.randn_like
|
|
def capture_noise(*noise_args, **noise_kwargs):
|
|
noise = original_randn_like(*noise_args, **noise_kwargs)
|
|
captured["vae.encoder_noise"] = noise.detach().cpu().contiguous().clone()
|
|
return noise
|
|
torch.randn_like = capture_noise
|
|
try:
|
|
output, lengths = original_encoding(*encoding_args, **encoding_kwargs)
|
|
finally:
|
|
torch.randn_like = original_randn_like
|
|
captured["vae.reference_latents"] = output.detach().cpu().contiguous().clone()
|
|
captured["vae.reference_lengths"] = lengths.detach().cpu().contiguous().clone()
|
|
return output, lengths
|
|
engine.vae_model.encoding_and_normalization = capture_encoding
|
|
def capture_audio_encoder(_module, _inputs, output):
|
|
captured["qwen.audio_output"] = output.last_hidden_state.detach().cpu().contiguous().clone()
|
|
handles.append(engine.model.text_encoder.audio_tower.register_forward_pre_hook(
|
|
capture_inputs("qwen.audio"), with_kwargs=True))
|
|
handles.append(engine.model.text_encoder.audio_tower.register_forward_hook(capture_audio_encoder))
|
|
def capture_audio_boundary(name):
|
|
def hook(_module, inputs, output):
|
|
captured[f"qwen.audio.{name}.input"] = inputs[0].detach().cpu().contiguous().clone()
|
|
value = output[0] if isinstance(output, tuple) else output
|
|
captured[f"qwen.audio.{name}.output"] = value.detach().cpu().contiguous().clone()
|
|
return hook
|
|
audio_tower = engine.model.text_encoder.audio_tower
|
|
for name in ("conv1", "conv2", "ln_post", "proj"):
|
|
handles.append(getattr(audio_tower, name).register_forward_hook(capture_audio_boundary(name)))
|
|
for index in (0, len(audio_tower.layers) - 1):
|
|
handles.append(audio_tower.layers[index].register_forward_hook(capture_audio_boundary(f"layer.{index}")))
|
|
captured["qwen.audio.positional_embedding"] = (
|
|
audio_tower.positional_embedding.positional_embedding.detach().cpu().contiguous().clone())
|
|
|
|
def capture_decoder_input(_module, inputs):
|
|
captured["vae.input"] = inputs[0].detach().cpu().contiguous().clone()
|
|
|
|
handles.append(engine.vae_model.conv_pre.register_forward_pre_hook(capture_decoder_input))
|
|
content = [{"type": "text", "text": args.instruction}]
|
|
if args.audio:
|
|
content.append({"type": "audio", "audio": str(args.audio.resolve())})
|
|
messages = [{"role": "user", "content": content}]
|
|
# Initialization consumes RNG. Reset before VAE reference sampling as well
|
|
# as the diffusion seed set inside the reference sample implementation.
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
try:
|
|
waveform, sample_rate = engine.generate(
|
|
messages,
|
|
gen_seconds=args.duration,
|
|
nfe=args.steps,
|
|
cfg_strength=args.cfg,
|
|
sway_sampling_coef=args.sway,
|
|
seed=args.seed,
|
|
)
|
|
finally:
|
|
engine.model.sample = original_sample
|
|
engine.vae_model.denormalize = original_denormalize
|
|
engine.vae_model.encoding_and_normalization = original_encoding
|
|
for handle in handles:
|
|
handle.remove()
|
|
|
|
captured["waveform"] = waveform.contiguous()
|
|
save_file(captured, args.output_dir / "reference.safetensors")
|
|
if args.qkv_fixture:
|
|
for name in ("q_proj", "k_proj", "v_proj"):
|
|
if not torch.equal(captured[f"qwen.layer.{name}"].float(), qkv_fixture[name]):
|
|
raise RuntimeError(f"Replayed projection differs: {name}")
|
|
actual = qkv_fixture["attention_output"].double().flatten()
|
|
expected = captured["qwen.attention_output"].double().flatten()
|
|
cosine = torch.nn.functional.cosine_similarity(actual, expected, dim=0).item()
|
|
rmse = (actual - expected).square().mean().sqrt().item()
|
|
logging.info("Matched QKV attention cosine=%.12f rmse=%.12g", cosine, rmse)
|
|
save_audio(waveform, sample_rate, str(args.output_dir / "reference.wav"))
|
|
manifest = {
|
|
"instruction": args.instruction,
|
|
"messages": messages,
|
|
"seed": args.seed,
|
|
"duration": args.duration,
|
|
"steps": args.steps,
|
|
"cfg": args.cfg,
|
|
"sway": args.sway,
|
|
"dtype": args.dtype,
|
|
"qkv_replay_fixture": str(args.qkv_fixture) if args.qkv_fixture else None,
|
|
"torch": torch.__version__,
|
|
"tf32": False,
|
|
"bf16_reduced_precision_reduction": torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction,
|
|
"sample_rate": sample_rate,
|
|
"tensors": {key: {"shape": list(value.shape), "dtype": str(value.dtype)}
|
|
for key, value in captured.items()},
|
|
"performance_measurement": False,
|
|
}
|
|
(args.output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
|
|
logging.info("Saved reference fixture: %s", args.output_dir)
|
|
if args.benchmark_runs:
|
|
milliseconds = []
|
|
expected_waveform = waveform.cpu()
|
|
for run in range(args.benchmark_runs + 2):
|
|
if args.audio:
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
repeated, repeated_rate = engine.generate(messages, gen_seconds=args.duration,
|
|
nfe=args.steps, cfg_strength=args.cfg, sway_sampling_coef=args.sway, seed=args.seed)
|
|
repeated = repeated.cpu()
|
|
torch.cuda.synchronize()
|
|
elapsed = (time.perf_counter() - start) * 1000
|
|
if repeated_rate != sample_rate or not torch.equal(repeated, expected_waveform):
|
|
raise RuntimeError("benchmark waveform changed")
|
|
if run >= 2:
|
|
milliseconds.append(elapsed)
|
|
logging.info("AuK inference_ms=%.6f", elapsed)
|
|
result = {"runs": args.benchmark_runs, "milliseconds": milliseconds,
|
|
"median_ms": statistics.median(milliseconds), "dtype": args.dtype,
|
|
"rtf": statistics.median(milliseconds) / (1000 * args.duration),
|
|
"scope": "loaded original Python generate to CPU waveform; capture hooks removed; excludes WAV writing"}
|
|
(args.output_dir / "inference-performance.json").write_text(json.dumps(result, indent=2) + "\n")
|
|
logging.info("AuK inference_median_ms=%.6f RTF=%.6f", result["median_ms"], result["rtf"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|