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
94 lines
4.6 KiB
Python
94 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Run Sidon's released TorchScript pipeline and capture component boundaries."""
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import time
|
|
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
import torchaudio
|
|
from safetensors.torch import save_file
|
|
from transformers import SeamlessM4TFeatureExtractor
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-dir", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--parity", action="store_true")
|
|
parser.add_argument("--warmup", type=int, default=2)
|
|
parser.add_argument("--repeat", type=int, default=3)
|
|
args = parser.parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
if args.parity:
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
encoder = torch.jit.load(str(args.model_dir / "feature_extractor_cuda.pt"), map_location="cuda").eval()
|
|
decoder = torch.jit.load(str(args.model_dir / "decoder_cuda.pt"), map_location="cuda").eval()
|
|
# facebook/w2v-bert-2.0/preprocessor_config.json differs from the class default.
|
|
frontend = SeamlessM4TFeatureExtractor(padding_value=1)
|
|
original, sample_rate = sf.read(args.audio, dtype="float64", always_2d=True)
|
|
peak = np.abs(original).max()
|
|
if peak == 0:
|
|
raise ValueError("The official peak normalization is undefined for silent input")
|
|
waveform = torch.from_numpy((0.9 * (original / peak)).astype(np.float32)).mean(dim=1).view(1, -1)
|
|
original_frames = original.shape[0]
|
|
target_samples = int(48000 / sample_rate * original_frames)
|
|
measurements = []
|
|
with torch.inference_mode():
|
|
for run in range(args.warmup + args.repeat):
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
start = time.perf_counter()
|
|
wav = torchaudio.functional.highpass_biquad(waveform, sample_rate, 50)
|
|
wav = torchaudio.functional.resample(wav, sample_rate, 16000)
|
|
wav = torch.nn.functional.pad(wav, (0, 24000))
|
|
restored = []
|
|
cache = None
|
|
stages = {"frontend_ms": 0.0, "encoder_ms": 0.0, "decoder_ms": 0.0}
|
|
fixtures = {}
|
|
for index, chunk in enumerate(wav.view(-1).split(16000 * 96)):
|
|
stage = time.perf_counter()
|
|
padded = torch.nn.functional.pad(chunk, (160, 160))
|
|
features = frontend(padded, sampling_rate=16000, return_tensors="pt")["input_features"].cuda()
|
|
torch.cuda.synchronize()
|
|
stages["frontend_ms"] += (time.perf_counter() - stage) * 1000
|
|
stage = time.perf_counter()
|
|
hidden = encoder(features)["last_hidden_state"]
|
|
torch.cuda.synchronize()
|
|
stages["encoder_ms"] += (time.perf_counter() - stage) * 1000
|
|
if cache is not None:
|
|
hidden = torch.cat([cache, hidden], dim=1)
|
|
stage = time.perf_counter()
|
|
audio = decoder(hidden.transpose(1, 2)).view(-1)[:-960]
|
|
torch.cuda.synchronize()
|
|
stages["decoder_ms"] += (time.perf_counter() - stage) * 1000
|
|
restored.append(audio)
|
|
cache = hidden[:, -1:]
|
|
if args.parity:
|
|
fixtures[f"chunk{index}.padded"] = padded.contiguous()
|
|
fixtures[f"chunk{index}.features"] = features.cpu().contiguous()
|
|
fixtures[f"chunk{index}.hidden"] = hidden.cpu().contiguous()
|
|
fixtures[f"chunk{index}.audio"] = audio.cpu().contiguous()
|
|
output = torch.cat(restored)[:target_samples].cpu().numpy()
|
|
elapsed = (time.perf_counter() - start) * 1000
|
|
row = {"warmup": run < args.warmup, "session.wall_ms": elapsed,
|
|
"audio_seconds": original_frames / sample_rate,
|
|
"rtf": elapsed / 1000 / (original_frames / sample_rate),
|
|
"peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20,
|
|
"peak_reserved_mib": torch.cuda.max_memory_reserved() / 2**20,
|
|
**stages}
|
|
print(json.dumps(row), flush=True)
|
|
measurements.append(row)
|
|
sf.write(args.output_dir / "audio.wav", output, 48000, subtype="FLOAT")
|
|
if args.parity:
|
|
save_file(fixtures, args.output_dir / "components.safetensors")
|
|
(args.output_dir / "measurements.json").write_text(json.dumps(measurements, indent=2) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|