File size: 3,467 Bytes
dd51636 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | #!/usr/bin/env python
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import torch
import torchaudio
import onnxruntime as ort
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Decode saved decoder[4] features with exported iSTFTNet2.")
parser.add_argument("--feature", required=True, help="Path to .npy feature array with shape [frames, 768].")
parser.add_argument("--decoder-dir", default="istftnet2_decoder4_50hz")
parser.add_argument("--runtime", default="torchscript_cuda")
parser.add_argument("--output", default="outputs/decoder4_feature_48k.wav")
return parser.parse_args()
def run_torchscript(features: torch.Tensor, decoder_dir: Path, runtime: str) -> torch.Tensor:
if runtime == "torchscript_cuda":
if not torch.cuda.is_available():
raise RuntimeError("torchscript_cuda requested, but torch.cuda.is_available() is false.")
artifact = decoder_dir / "istftnet2_decoder_cuda.ts"
device = torch.device("cuda")
elif runtime == "torchscript_cpu":
artifact = decoder_dir / "istftnet2_decoder_cpu.ts"
device = torch.device("cpu")
else:
raise ValueError(runtime)
decoder = torch.jit.load(str(artifact), map_location=device).eval()
with torch.inference_mode():
return decoder(features.to(device)).detach().float().cpu()
def run_onnx(features: torch.Tensor, decoder_dir: Path, runtime: str) -> torch.Tensor:
provider_map = {
"onnx_rocm": "ROCMExecutionProvider",
"onnx_cuda": "CUDAExecutionProvider",
"onnx_cpu": "CPUExecutionProvider",
}
provider = provider_map[runtime]
available = ort.get_available_providers()
if provider not in available:
raise RuntimeError(f"{provider} is unavailable. Available providers: {available}")
session = ort.InferenceSession(str(decoder_dir / "istftnet2_decoder.onnx"), providers=[provider])
input_name = session.get_inputs()[0].name
audio = session.run(None, {input_name: features.numpy().astype(np.float32, copy=False)})[0]
return torch.from_numpy(audio).float()
def main() -> None:
args = parse_args()
decoder_dir = Path(args.decoder_dir)
arr = np.load(args.feature).astype(np.float32, copy=False)
if arr.ndim != 2 or arr.shape[1] != 768:
raise ValueError(f"Expected feature shape [frames, 768], got {arr.shape}.")
features = torch.from_numpy(arr).T.unsqueeze(0).contiguous()
if args.runtime.startswith("torchscript_"):
audio = run_torchscript(features, decoder_dir, args.runtime)
else:
audio = run_onnx(features, decoder_dir, args.runtime)
wav = audio[0].clamp(-1.0, 1.0).contiguous()
expected_samples = int(arr.shape[0]) * 960
wav = wav[..., :expected_samples]
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
torchaudio.save(str(output), wav, 48000)
manifest = {
"output": str(output),
"feature": args.feature,
"decoder_dir": args.decoder_dir,
"runtime": args.runtime,
"sample_rate": 48000,
"feature_shape": list(arr.shape),
"num_samples": int(wav.shape[-1]),
"duration_sec": float(wav.shape[-1] / 48000),
}
output.with_suffix(".json").write_text(json.dumps(manifest, indent=2))
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()
|