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.6 KiB
Python
60 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Save official MusicCoCa outputs for the independent C++ parity test."""
|
|
|
|
import argparse
|
|
import logging
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import numpy as np
|
|
from safetensors.numpy import save_file
|
|
import soundfile as sf
|
|
from scipy.signal import resample_poly
|
|
from threadpoolctl import threadpool_limits
|
|
|
|
from magenta_rt import musiccoca
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--resources", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--log", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(args.log), logging.StreamHandler()], force=True)
|
|
model = musiccoca.MusicCoCa(args.resources)
|
|
outputs = {}
|
|
prompts = ("disco funk", "A gentle acoustic guitar with warm piano", "Driving electronic drums and deep bass")
|
|
noise = np.random.RandomState(0).randn(768).astype(np.float32)
|
|
outputs["mapper_noise"] = noise
|
|
with threadpool_limits(limits=8):
|
|
for i, prompt in enumerate(prompts):
|
|
start = time.perf_counter()
|
|
text = model.embed(prompt, use_mapper=False)
|
|
mapped = model.embed(prompt, use_mapper=True, seed=0)
|
|
codes = model.tokenize(mapped).astype(np.int32)
|
|
outputs[f"text_{i}"] = text.reshape(-1)
|
|
outputs[f"mapped_{i}"] = mapped.reshape(-1)
|
|
outputs[f"codes_{i}"] = codes.reshape(-1)
|
|
logging.info("prompt=%r session.wall_ms=%.3f", prompt, (time.perf_counter()-start)*1000)
|
|
audio, sample_rate = sf.read(args.audio, dtype="float32", always_2d=True)
|
|
audio = audio.mean(axis=1)
|
|
divisor = np.gcd(sample_rate, 16000)
|
|
audio = resample_poly(audio, 16000 // divisor, sample_rate // divisor).astype(np.float32)
|
|
audio = np.pad(audio[:160000], (0, max(0, 160000-len(audio))))
|
|
prep = model._audio_preprocessor
|
|
prep.set_tensor(prep.get_input_details()[0]["index"], audio.reshape(1, -1))
|
|
prep.invoke()
|
|
features = prep.get_tensor(prep.get_output_details()[0]["index"])
|
|
encoder = model._music_encoder
|
|
encoder.set_tensor(encoder.get_input_details()[0]["index"], features)
|
|
encoder.invoke()
|
|
outputs["audio_features"] = features.reshape(-1)
|
|
outputs["audio_input"] = audio
|
|
outputs["audio_embedding"] = encoder.get_tensor(encoder.get_output_details()[0]["index"]).reshape(-1)
|
|
save_file(outputs, args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|