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
60 lines
2.7 KiB
Python
60 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Measure the unmodified official Python system using the C++ validation cases."""
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import jax
|
|
import soundfile as sf
|
|
from ai_edge_litert.interpreter import Interpreter
|
|
|
|
from magenta_rt import paths, musiccoca
|
|
from magenta_rt.config import MUSICCOCA
|
|
from magenta_rt.jax.system import MagentaRT2System
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--source", type=Path, required=True)
|
|
parser.add_argument("--variant", choices=("small", "base"), required=True)
|
|
parser.add_argument("--cases", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--log", type=Path, required=True)
|
|
parser.add_argument("--threads", type=int, default=8)
|
|
args = parser.parse_args()
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(level=logging.INFO, handlers=[logging.StreamHandler(), logging.FileHandler(args.log)], force=True)
|
|
paths.set_magenta_home(args.source)
|
|
def interpreter_with_threads(model_path):
|
|
interpreter = Interpreter(model_path=str(model_path), num_threads=args.threads)
|
|
interpreter.allocate_tensors()
|
|
return interpreter
|
|
musiccoca._make_interpreter = interpreter_with_threads
|
|
model = MagentaRT2System(size="mrt2_" + args.variant)
|
|
cases = json.loads(args.cases.read_text())
|
|
with (args.output / "metrics.csv").open("w") as output:
|
|
writer = csv.DictWriter(output, fieldnames=("case", "audio_sec", "wall_ms", "rtf"))
|
|
writer.writeheader()
|
|
for case in cases:
|
|
if case["seed"] != 0:
|
|
raise ValueError("The upstream Python system fixes its generation RNG seed to zero")
|
|
jax.block_until_ready(model._params)
|
|
start = time.perf_counter()
|
|
embedding = model.embed_style(case["text"], use_mapper=True)
|
|
# Keep generation on the official, combined AR+codec streaming function.
|
|
waveform, _ = model.generate(conditioning={MUSICCOCA.key: embedding}, frames=round(case["duration_sec"] * 25))
|
|
wall_ms = (time.perf_counter() - start) * 1000
|
|
duration = len(waveform.samples) / waveform.sample_rate
|
|
sf.write(args.output / (case["id"] + ".wav"), waveform.samples, waveform.sample_rate, subtype="PCM_16")
|
|
row = {"case": case["id"], "audio_sec": duration, "wall_ms": wall_ms, "rtf": wall_ms / (1000 * duration)}
|
|
writer.writerow(row)
|
|
output.flush()
|
|
logging.info("case=%s session.wall_ms=%.3f audio.duration_sec=%.3f RTF=%.6f", case["id"], wall_ms, duration, row["rtf"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|