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
79 lines
3.6 KiB
Python
79 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate original Transformers Bark 3.0 Pro references from local WAV files.
|
|
|
|
Use the pinned source snapshot described in docs/models/underdog_bark3.md.
|
|
This deliberately does not import or use audio.cpp's prepared weights.
|
|
"""
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from transformers import AutoProcessor, ParakeetForTDT
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, nargs="+", required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
|
|
parser.add_argument("--dtype", choices=["float32", "bfloat16"], default="float32")
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--warmup-audio", type=Path)
|
|
args = parser.parse_args()
|
|
if args.iterations < 1 or args.threads < 1:
|
|
parser.error("iterations and threads must be positive")
|
|
if args.device == "cpu" and args.dtype != "float32":
|
|
parser.error("the publisher's CPU fallback uses float32")
|
|
if args.device == "cuda" and not torch.cuda.is_available():
|
|
parser.error("CUDA is unavailable in this PyTorch environment")
|
|
if args.device == "cuda" and args.dtype == "bfloat16" and not torch.cuda.is_bf16_supported():
|
|
parser.error("this GPU does not support BF16")
|
|
torch.set_num_threads(args.threads)
|
|
dtype = getattr(torch, args.dtype)
|
|
processor = AutoProcessor.from_pretrained(args.source, local_files_only=True)
|
|
model = ParakeetForTDT.from_pretrained(args.source, dtype=dtype, local_files_only=True).to(args.device).eval()
|
|
def read_wave(path):
|
|
wave, rate = sf.read(path, dtype="float32")
|
|
if rate != 16000 or wave.ndim != 1 or len(wave) == 0:
|
|
raise ValueError("reference fixtures must be nonempty mono 16 kHz WAVs")
|
|
return wave
|
|
def inputs_for(wave):
|
|
# Same once-per-recording prefix as runtime/transcribe_bark_pro.py.
|
|
padded = np.concatenate([np.zeros(4000, np.float32), wave])
|
|
return processor(padded, sampling_rate=16000, return_tensors="pt").to(args.device, dtype)
|
|
def synchronize():
|
|
if args.device == "cuda":
|
|
torch.cuda.synchronize()
|
|
if args.warmup_audio:
|
|
with torch.inference_mode():
|
|
model.generate(**inputs_for(read_wave(args.warmup_audio)))
|
|
synchronize()
|
|
records = []
|
|
for path in args.audio:
|
|
wave = read_wave(path) # File I/O and model loading are outside the timer.
|
|
runs = []
|
|
for _ in range(args.iterations):
|
|
synchronize()
|
|
started = time.perf_counter()
|
|
inputs = inputs_for(wave)
|
|
with torch.inference_mode():
|
|
generated = model.generate(**inputs)
|
|
synchronize()
|
|
ids = generated.sequences if hasattr(generated, "sequences") else generated
|
|
text = processor.batch_decode(ids, skip_special_tokens=True)[0]
|
|
runs.append(dict(seconds=time.perf_counter() - started, text=text))
|
|
records.append(dict(name=path.stem, path=str(path.resolve()), audio_seconds=len(wave) / 16000,
|
|
device=args.device, dtype=args.dtype, text=text, ids=ids.cpu().tolist(), runs=runs))
|
|
print(path.name, repr(text), flush=True)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(records, indent=2), encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|