File size: 1,610 Bytes
fb41b93 | 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 | import argparse
import json
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from chatterbox_s3gen_onestep_sdk.inference import ModelSession
from chatterbox_s3gen_onestep_sdk.postprocess import postprocess
from chatterbox_s3gen_onestep_sdk.preprocess import preprocess
def main():
parser = argparse.ArgumentParser(description="chatterbox-s3gen-onestep inference example")
parser.add_argument("--model", required=True, help="AXMODEL 路径")
parser.add_argument("--input", nargs="+", required=True, help="输入 npy（每输入一个，与 model_meta 顺序一致）")
parser.add_argument("--output-dir", default="output", help="输出目录")
args = parser.parse_args()
arrays = [np.load(p).astype(np.float32) for p in args.input]
session = ModelSession(args.model)
feeds = preprocess(*arrays)
raw = session.run_named(feeds)
result = postprocess(*raw)
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
for i, arr in enumerate(raw):
np.save(out_dir / f"output_{i}.npy", np.asarray(arr, dtype=np.float32))
try:
json.dumps(result)
(out_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
except TypeError:
np.save(out_dir / "result.npy", np.asarray(result, dtype=np.float32))
print("backend:", session.backend)
print("inputs:", session.input_names)
print("outputs:", session.output_names)
print("saved to:", out_dir)
if __name__ == "__main__":
main()
|