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
93 lines
4 KiB
Python
93 lines
4 KiB
Python
#!/usr/bin/env python3
|
|
"""Run pinned Plus's unmodified official loader against local mono 16 kHz WAVs.
|
|
|
|
Download bark_loader.py from the pinned Plus revision as well as the data/base
|
|
sidecars in the model guide. This never imports audio.cpp's converted weights.
|
|
"""
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
from transformers import AutoProcessor
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools/community_models"))
|
|
from convert_underdog_bark3_plus import BASE_SHA256, PLUS_PROFILE, sha256
|
|
|
|
LOADER_SHA256 = "e86b7f130bec1a8fa14dd315cfce95785b890dbd30000f6ceb2597fdd4312b02"
|
|
|
|
|
|
def main(profile=PLUS_PROFILE, description=__doc__):
|
|
parser = argparse.ArgumentParser(description=description)
|
|
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("--device", choices=["cpu", "cuda"], default="cpu")
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
parser.add_argument("--iterations", type=int, default=1)
|
|
parser.add_argument("--warmup-audio", type=Path)
|
|
parser.add_argument("--check-official-transcribe", action="store_true")
|
|
args = parser.parse_args()
|
|
if args.threads <= 0 or args.iterations <= 0:
|
|
parser.error("threads and iterations must be positive")
|
|
if args.device == "cuda" and not torch.cuda.is_available():
|
|
parser.error("CUDA is unavailable")
|
|
for name, expected in {**BASE_SHA256, "bark_loader.py": LOADER_SHA256, "model.bark": profile.source_sha256}.items():
|
|
if sha256(args.source / name) != expected:
|
|
raise ValueError("source differs from pinned snapshot: " + name)
|
|
torch.set_num_threads(args.threads)
|
|
spec = importlib.util.spec_from_file_location("official_bark_plus", args.source / "bark_loader.py")
|
|
official = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(official)
|
|
model = official.load_model(args.source / "model.bark", device=args.device, dtype=torch.float32)
|
|
processor = AutoProcessor.from_pretrained(args.source, local_files_only=True)
|
|
|
|
def read(path):
|
|
wave, rate = sf.read(path, dtype="float32")
|
|
if rate != 16000 or wave.ndim != 1 or not len(wave):
|
|
raise ValueError("reference inputs must be nonempty mono 16 kHz WAVs")
|
|
return wave
|
|
|
|
def sync():
|
|
if args.device == "cuda":
|
|
torch.cuda.synchronize()
|
|
|
|
def generate(wave):
|
|
padded = np.concatenate([np.zeros(4000, np.float32), wave])
|
|
inputs = processor([padded], sampling_rate=16000, padding=True, return_tensors="pt").to(args.device)
|
|
with torch.inference_mode():
|
|
result = model.generate(**inputs)
|
|
ids = result.sequences
|
|
return ids, processor.batch_decode(ids, skip_special_tokens=True)[0].strip()
|
|
|
|
if args.warmup_audio:
|
|
generate(read(args.warmup_audio))
|
|
sync()
|
|
records = []
|
|
for index, path in enumerate(args.audio):
|
|
wave = read(path)
|
|
runs = []
|
|
for _ in range(args.iterations):
|
|
sync()
|
|
start = time.perf_counter()
|
|
ids, text = generate(wave)
|
|
sync()
|
|
runs.append(dict(seconds=time.perf_counter() - start, text=text))
|
|
if args.check_official_transcribe and index == 0:
|
|
actual = official.transcribe(model, wave, 16000)
|
|
if actual != text:
|
|
raise AssertionError("pinned local processor differs from official transcribe()")
|
|
records.append(dict(name=path.stem, path=str(path.resolve()), audio_seconds=len(wave) / 16000,
|
|
device=args.device, dtype="float32", 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()
|