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
188 lines
9.7 KiB
Python
188 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the official Talker; capture only its AR input and final latent/waveform boundaries."""
|
|
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools/audiocpp_cli"))
|
|
from run_audiocpp_cli_path_tests import ResourceMonitor
|
|
|
|
|
|
class TalkerCapture:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.tensors = {}
|
|
self.prompts = []
|
|
self.encoder_layer_calls = []
|
|
self.original = model.generate_latents
|
|
self.decode = model.audio_vae.decode
|
|
self.inputs = model.input_embeddings
|
|
|
|
def __enter__(self):
|
|
self.encoder_hooks = [layer.register_forward_hook(
|
|
lambda module, inputs, output, index=index: self.encoder_layer_calls.append(index))
|
|
for index, layer in enumerate(self.model.audio_vae.encoder.encoder.layers)]
|
|
self.encoder_boundary_hook = self.model.audio_vae.encoder.register_forward_hook(
|
|
lambda module, inputs, output: self.tensors.update({
|
|
"reference.encoder_parameters": output[0].detach().cpu().contiguous(),
|
|
"reference.waveform": output[1].detach().cpu().contiguous()}))
|
|
self.speaker_hook = self.model.spk_head.register_forward_pre_hook(
|
|
lambda module, inputs: self.tensors.update({
|
|
"reference.speaker_features": inputs[0].detach().cpu().contiguous()}))
|
|
|
|
def prompt(*args, **kwargs):
|
|
name = f"chunk.{len(self.prompts)}"
|
|
handle = self.model.model.get_input_embeddings().register_forward_pre_hook(
|
|
lambda module, inputs: self.tensors.update({
|
|
name + ".prompt_ids": inputs[0].detach().to(torch.int32).cpu().contiguous()}))
|
|
try:
|
|
result = self.inputs(*args, **kwargs)
|
|
finally:
|
|
handle.remove()
|
|
self.prompts.append({"text": args[0], "user_prompt": args[1], "instruction": args[2],
|
|
"reference_text": args[4]})
|
|
return result
|
|
|
|
def generate(embeddings, reference, max_steps, cfg, sigma, temperature):
|
|
name = f"chunk.{len(self.prompts) - 1}"
|
|
offset_words = torch.cuda.default_generators[0].get_offset()
|
|
self.tensors[name + ".initial_rng_offset_words"] = torch.tensor(offset_words, dtype=torch.int64)
|
|
self.tensors[name + ".initial_rng_offset_blocks"] = torch.tensor(offset_words // 4, dtype=torch.int64)
|
|
self.tensors[name + ".prompt_embeddings"] = embeddings.detach().cpu().contiguous()
|
|
if reference is not None:
|
|
self.tensors[name + ".reference_latents"] = reference.detach().cpu().contiguous()
|
|
result = self.original(embeddings, reference, max_steps, cfg, sigma, temperature)
|
|
self.tensors[name + ".generated_latents"] = result.detach().cpu().contiguous()
|
|
return result
|
|
|
|
def decode(latents, *args, **kwargs):
|
|
result = self.decode(latents, *args, **kwargs)
|
|
self.tensors[f"chunk.{len(self.prompts) - 1}.waveform"] = result[0].detach().cpu().contiguous()
|
|
return result
|
|
|
|
self.model.input_embeddings = prompt
|
|
self.model.generate_latents = generate
|
|
self.model.audio_vae.decode = decode
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
for handle in self.encoder_hooks:
|
|
handle.remove()
|
|
self.encoder_boundary_hook.remove()
|
|
self.speaker_hook.remove()
|
|
self.model.generate_latents = self.original
|
|
self.model.audio_vae.decode = self.decode
|
|
self.model.input_embeddings = self.inputs
|
|
|
|
|
|
@torch.inference_mode()
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--requests", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--capture", action="store_true")
|
|
parser.add_argument("--decode-fixture", type=Path,
|
|
help="Teacher-force captured AR output through the official VAE only")
|
|
parser.add_argument("--vae-attention", choices=("eager", "sdpa"), default="eager")
|
|
parser.add_argument("--log", action="store_true", required=True)
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(args.model_dir.resolve()))
|
|
from talker import MingTalker
|
|
|
|
model = MingTalker.from_pretrained(args.model_dir / "talker", dtype=torch.bfloat16, device="cuda:0")
|
|
model.audio_vae.decoder.decoder.set_attn_implementation(args.vae_attention)
|
|
if args.capture and args.decode_fixture:
|
|
model.audio_vae.decoder.upsampling.upsampler.register_forward_hook(
|
|
lambda module, inputs, output: print(json.dumps({"vae.upsampler.input_dtype": str(inputs[0].dtype),
|
|
"vae.upsampler.output_dtype": str(output.dtype)}), flush=True))
|
|
model.audio_vae.decoder.decoder.register_forward_pre_hook(
|
|
lambda module, inputs, kwargs: print(json.dumps({"vae.decoder.input_dtype": str(kwargs["inputs_embeds"].dtype),
|
|
"vae.rope.inv_freq_dtype": str(module.rotary_emb.inv_freq.dtype)}), flush=True),
|
|
with_kwargs=True)
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
manifest = {"torch": torch.__version__, "gpu": torch.cuda.get_device_name(0), "requests": []}
|
|
requests = json.loads(args.requests.read_text(encoding="utf-8"))
|
|
if isinstance(requests, dict):
|
|
requests = requests["requests"]
|
|
for request in requests:
|
|
request = {**request, **request.get("options", {})}
|
|
name = request.get("name", request.get("id"))
|
|
if not name:
|
|
raise ValueError("Each request requires a name or id")
|
|
reference = None
|
|
reference_audio = request.get("reference_audio", request.get("voice_ref"))
|
|
if reference_audio:
|
|
samples, sample_rate = sf.read(reference_audio, dtype="float32", always_2d=True)
|
|
reference = {"samples": np.mean(samples, axis=1), "sample_rate": sample_rate}
|
|
monitor = ResourceMonitor(os.getpid(), 0, 20)
|
|
stop = threading.Event()
|
|
|
|
def sample_resources():
|
|
while not stop.is_set():
|
|
monitor.sample()
|
|
stop.wait(0.02)
|
|
|
|
capture = TalkerCapture(model)
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
job = executor.submit(sample_resources)
|
|
try:
|
|
start = time.perf_counter()
|
|
if args.capture and not args.decode_fixture:
|
|
capture.__enter__()
|
|
if args.decode_fixture:
|
|
tensors = load_file(args.decode_fixture, device="cuda:0")
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
waveform = model.audio_vae.decode(tensors["chunk.0.generated_latents"], use_cache=False, last_chunk=True)[0]
|
|
result = {"audio": waveform.detach().float().cpu(), "sr": model.audio_vae.config.sample_rate}
|
|
else:
|
|
result = model.synthesize(
|
|
request["text"], ip_name=request.get("voice_id", "US_James"), reference=reference,
|
|
prompt_text=request.get("reference_text"), language=request.get("language"),
|
|
max_steps=request.get("max_audio_tokens", 1000), seed=request.get("seed", 42),
|
|
cfg=request.get("guidance_scale", 2), sigma=request.get("sigma", 0.25),
|
|
temperature=request.get("temperature", 0))
|
|
torch.cuda.synchronize()
|
|
wall_ms = (time.perf_counter() - start) * 1000
|
|
finally:
|
|
if args.capture and not args.decode_fixture:
|
|
capture.__exit__()
|
|
stop.set()
|
|
job.result()
|
|
monitor.close()
|
|
waveform = result["audio"].reshape(-1).numpy()
|
|
sample_rate = int(result["sr"])
|
|
sf.write(args.output_dir / f"{name}.wav", waveform, sample_rate, subtype="FLOAT")
|
|
if args.capture and not args.decode_fixture:
|
|
save_file(capture.tensors, str(args.output_dir / f"{name}.safetensors"))
|
|
duration = len(waveform) / sample_rate
|
|
record = {**request, "session.wall_ms": wall_ms, "audio_duration_sec": duration,
|
|
"rtf": wall_ms / (duration * 1000), "instrumented": args.capture,
|
|
"peak_allocated_vram_bytes": torch.cuda.max_memory_allocated(),
|
|
"peak_reserved_vram_bytes": torch.cuda.max_memory_reserved(),
|
|
**monitor.metrics.to_json(), "decode_only": bool(args.decode_fixture),
|
|
"vae_attention": args.vae_attention, "prompts": capture.prompts,
|
|
"reference_encoder_modules": len(model.audio_vae.encoder.encoder.layers),
|
|
"reference_encoder_config_layers": model.audio_vae.encoder.encoder.config.num_hidden_layers,
|
|
"reference_encoder_shared_config": model.audio_vae.encoder.encoder.config is model.audio_vae.encoder.aggregator.config,
|
|
"reference_encoder_executed_layers": capture.encoder_layer_calls}
|
|
manifest["requests"].append(record)
|
|
print(json.dumps(record, ensure_ascii=False), flush=True)
|
|
(args.output_dir / "metrics.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|