File size: 2,349 Bytes
0d06934 | 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 | """Convert the Haas-ja-v1.0 FastConformer encoder to CoreML (.mlpackage, fp16).
Usage: python convert.py <T_mel> [out.mlpackage]
Input: audio_signal [1,128,T_mel] f32 (zero past length), length [1] i32 -> encoded [1,1024,T_mel/8]
"""
import os, sys
# HAAS_EXPORT_DIR: dir containing the haas_onnx package (scripts/ of Sanami/haas-ja-onnx)
# HAAS_NEMO_DIR: the extracted Haas-ja-v1.0.nemo (model_config.yaml, model_weights.ckpt)
sys.path.insert(0, os.environ.get("HAAS_EXPORT_DIR", "export"))
import numpy as np, torch, coremltools as ct
from haas_onnx.load import load_model
import cmlpatch # noqa: F401
T = int(sys.argv[1])
out = sys.argv[2] if len(sys.argv) > 2 else f"encoder_T{T}.mlpackage"
class Enc(torch.nn.Module):
def __init__(self, enc):
super().__init__()
self.enc = enc
def forward(self, audio_signal, length=None):
if length is None:
length = torch.full((1,), audio_signal.shape[-1], dtype=torch.long)
y, _ = self.enc(audio_signal=audio_signal, length=length.long())
return y
model = load_model(os.environ.get("HAAS_NEMO_DIR", "nemo_x"))
enc = model.encoder.eval()
NL = int(os.environ.get("NL", "0")) # debug: keep only the first NL layers
if NL:
enc.layers = enc.layers[:NL]
NOLEN = os.environ.get("NOLEN") == "1" # debug: no length input (whole window valid)
enc.update_max_seq_length(T, device="cpu") if hasattr(enc, "update_max_seq_length") else None
w = Enc(enc).eval()
x = torch.randn(1, 128, T)
L = torch.tensor([T], dtype=torch.int32)
args = (x,) if NOLEN else (x, L)
with torch.no_grad():
ref = w(*args).numpy()
tr = torch.jit.trace(w, args, check_trace=False)
ins = [ct.TensorType("audio_signal", shape=(1, 128, T), dtype=np.float32)]
if not NOLEN:
ins.append(ct.TensorType("length", shape=(1,), dtype=np.int32))
ml = ct.convert(
tr,
inputs=ins,
outputs=[ct.TensorType("encoded", dtype=np.float32)],
convert_to="mlprogram", compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.iOS17,
)
ml.save(out)
feed = {"audio_signal": x.numpy()}
if not NOLEN:
feed["length"] = L.numpy()
y = ml.predict(feed)["encoded"]
print("shape", y.shape, "max abs diff", float(np.abs(y - ref).max()), "ref max", float(np.abs(ref).max()),
"cos", float((y * ref).sum() / np.linalg.norm(y) / np.linalg.norm(ref)))
|