#!/usr/bin/env python3 """ Unified CoreML conversion script for Parakeet EOU streaming encoder. Supports 160ms, 320ms, and 1600ms chunk sizes by properly configuring NeMo's streaming parameters before tracing. Usage: # 160ms (default) python convert_streaming_encoder_unified.py --chunk-ms 160 --output-dir Models/160ms/160ms # 320ms python convert_streaming_encoder_unified.py --chunk-ms 320 --output-dir Models/320ms # 1600ms python convert_streaming_encoder_unified.py --chunk-ms 1600 --output-dir Models/1600ms """ import argparse import json from pathlib import Path from typing import Tuple import coremltools as ct import numpy as np import torch import torch.nn as nn from nemo.collections.asr.models import EncDecRNNTBPEModel class LoopbackEncoderWrapper(nn.Module): """ Wraps the Parakeet Encoder for CoreML Loopback Streaming. This wrapper handles the pre_cache concatenation and cache management that NeMo does internally in its streaming pipeline. Inputs: - audio_signal: [B, D, T] (Mel spectrogram chunk) - audio_length: [B] - pre_cache: [B, D, pre_cache_size] (Previous mel context) - cache_last_channel: [layers, B, cache_size, hidden] - cache_last_time: [layers, B, hidden, time_cache] - cache_last_channel_len: [B] Outputs: - encoded_output: [B, D_out, T_out] - encoded_length: [B] - new_pre_cache: [B, D, pre_cache_size] - new_cache_last_channel - new_cache_last_time - new_cache_last_channel_len """ def __init__(self, encoder, pre_cache_size: int): super().__init__() self.encoder = encoder self.pre_cache_size = pre_cache_size def forward( self, audio_signal: torch.Tensor, audio_length: torch.Tensor, pre_cache: torch.Tensor, cache_last_channel: torch.Tensor, cache_last_time: torch.Tensor, cache_last_channel_len: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # 1. Prepend pre_cache to audio_signal full_input = torch.cat([pre_cache, audio_signal], dim=2) full_length = audio_length + self.pre_cache_size # 2. Extract NEW pre_cache (last N frames of full_input) new_pre_cache = full_input[:, :, -self.pre_cache_size :] # 3. Process with Encoder using cache_aware_stream_step encoded, encoded_len, new_cache_channel, new_cache_time, new_cache_len = ( self.encoder.cache_aware_stream_step( processed_signal=full_input, processed_signal_length=full_length, cache_last_channel=cache_last_channel, cache_last_time=cache_last_time, cache_last_channel_len=cache_last_channel_len, ) ) # Cast lengths to Int32 for CoreML encoded_len_32 = encoded_len.to(dtype=torch.int32) new_channel_len_32 = new_cache_len.to(dtype=torch.int32) return ( encoded, encoded_len_32, new_pre_cache, new_cache_channel, new_cache_time, new_channel_len_32, ) def get_streaming_config(encoder, chunk_ms: int): """ Get the correct streaming configuration for the given chunk size. Returns: dict with: - chunk_size: encoder output steps - shift_size: shift steps - mel_frames: input mel frames for the chunk - pre_cache_size: pre-encode cache size in mel frames - valid_out_len: number of valid output frames per chunk """ if chunk_ms == 160: # Default 160ms config - no need to call setup_streaming_params # Uses: chunk_size=[9, 16], valid_out_len=2 # With pre_cache=16, input is 17 mel frames return { "chunk_size": 4, # encoder output steps "shift_size": 2, "mel_frames": 17, # 16 + 1 for padding "pre_cache_size": 16, "valid_out_len": 2, "samples": 2560, # 160ms * 16000 } elif chunk_ms == 320: # 320ms config - need to call setup_streaming_params(chunk_size=8, shift_size=4) # After setup: chunk_size=[57, 64], shift_size=[25, 32], pre_encode_cache_size=[0, 9] # valid_out_len=4 encoder.setup_streaming_params(chunk_size=8, shift_size=4) cfg = encoder.streaming_cfg print(f"320ms streaming_cfg: {cfg}") # Use NeMo's exact values from streaming_cfg: # - chunk_size[1] = 64 mel frames # - shift_size[1] = 32 mel frames (320ms latency) # - pre_encode_cache_size[1] = 9 return { "chunk_size": 8, # encoder output steps "shift_size": 4, "mel_frames": 64, # From NeMo cfg.chunk_size[1] = 64 mel frames "pre_cache_size": 9, # From NeMo cfg.pre_encode_cache_size[1] = 9 "valid_out_len": 4, "samples": 10240, # 64 mel frames * 160 hop_length = 10240 samples (~640ms audio) "shift_samples": 5120, # 32 mel frames * 160 = 5120 samples (320ms latency) } elif chunk_ms == 1600: # 1600ms config - need to call setup_streaming_params(chunk_size=40, shift_size=20) # After setup: chunk_size=[313, 320] mel frames, shift_size=[153, 160] mel frames # The "1600ms" refers to the latency (shift), not the chunk size! # - Chunk size: 320 mel frames (3183ms of audio) # - Shift: 160 mel frames (1600ms latency) # - Output: valid_out_len=20 encoder frames encoder.setup_streaming_params(chunk_size=40, shift_size=20) cfg = encoder.streaming_cfg print(f"1600ms streaming_cfg: {cfg}") # Audio samples for 320 mel frames: (320-1)*160 + 400 - 512 = 50928 samples (~3183ms) # Audio shift for 160 mel frames: 160*160 = 25600 samples (1600ms) return { "chunk_size": 40, # encoder output steps "shift_size": 20, "mel_frames": 320, # From NeMo cfg.chunk_size[1] = 320 mel frames "pre_cache_size": 9, # From NeMo cfg.pre_encode_cache_size[1] = 9 "valid_out_len": 20, "samples": 50928, # (320-1)*160 + 400 - 512 = 50928 samples (~3183ms) "shift_samples": 25600, # 160*160 = 25600 samples (1600ms) } else: raise ValueError(f"Unsupported chunk size: {chunk_ms}ms. Use 160, 320, or 1600.") def main(): parser = argparse.ArgumentParser(description="Convert Parakeet EOU encoder to CoreML") parser.add_argument( "--chunk-ms", type=int, default=160, choices=[160, 320, 1600], help="Chunk size in milliseconds (160, 320, or 1600)", ) parser.add_argument( "--output-dir", type=str, default=None, help="Output directory for the CoreML model", ) parser.add_argument( "--model-id", type=str, default="nvidia/parakeet_realtime_eou_120m-v1", help="HuggingFace model ID", ) args = parser.parse_args() # Default output directory if args.output_dir is None: args.output_dir = f"Models/{args.chunk_ms}ms" output_path = Path(args.output_dir) output_path.mkdir(parents=True, exist_ok=True) print(f"Loading model: {args.model_id}...") asr_model = EncDecRNNTBPEModel.from_pretrained(model_name=args.model_id) asr_model.eval() encoder = asr_model.encoder # Get streaming configuration config = get_streaming_config(encoder, args.chunk_ms) mel_dim = 128 hidden_dim = encoder.d_model # 512 num_layers = len(encoder.layers) # 17 # Cache sizes cache_channel_size = 70 cache_time_size = 8 pre_cache_size = config["pre_cache_size"] chunk_size_in = config["mel_frames"] print(f"\n=== Configuration for {args.chunk_ms}ms ===") print(f"Mel frames: {chunk_size_in}") print(f"Pre-cache: {pre_cache_size}") print(f"Valid output len: {config['valid_out_len']}") print(f"Hidden dim: {hidden_dim}, Layers: {num_layers}") print(f"Cache: Channel={cache_channel_size}, Time={cache_time_size}") # Create wrapper wrapper = LoopbackEncoderWrapper(encoder, pre_cache_size=pre_cache_size) wrapper.eval() # Test inputs for tracing batch_size = 1 test_mel = torch.randn(batch_size, mel_dim, chunk_size_in) test_mel_len = torch.tensor([chunk_size_in], dtype=torch.int32) test_pre_cache = torch.zeros(batch_size, mel_dim, pre_cache_size) test_cache_channel = torch.zeros(num_layers, batch_size, cache_channel_size, hidden_dim) test_cache_time = torch.zeros(num_layers, batch_size, hidden_dim, cache_time_size) test_cache_len = torch.zeros(batch_size, dtype=torch.int32) print("\nTracing model...") traced_model = torch.jit.trace( wrapper, (test_mel, test_mel_len, test_pre_cache, test_cache_channel, test_cache_time, test_cache_len), strict=False, ) # Test output shape with torch.no_grad(): out = traced_model( test_mel, test_mel_len, test_pre_cache, test_cache_channel, test_cache_time, test_cache_len ) print(f"Encoder output shape: {out[0].shape}") # [B, D, T_out] # CoreML conversion print("\nConverting to CoreML...") inputs = [ ct.TensorType(name="audio_signal", shape=(1, mel_dim, chunk_size_in), dtype=np.float32), ct.TensorType(name="audio_length", shape=(1,), dtype=np.int32), ct.TensorType(name="pre_cache", shape=(1, mel_dim, pre_cache_size), dtype=np.float32), ct.TensorType( name="cache_last_channel", shape=(num_layers, 1, cache_channel_size, hidden_dim), dtype=np.float32, ), ct.TensorType( name="cache_last_time", shape=(num_layers, 1, hidden_dim, cache_time_size), dtype=np.float32, ), ct.TensorType(name="cache_last_channel_len", shape=(1,), dtype=np.int32), ] outputs = [ ct.TensorType(name="encoded_output", dtype=np.float32), ct.TensorType(name="encoded_length", dtype=np.int32), ct.TensorType(name="new_pre_cache", dtype=np.float32), ct.TensorType(name="new_cache_last_channel", dtype=np.float32), ct.TensorType(name="new_cache_last_time", dtype=np.float32), ct.TensorType(name="new_cache_last_channel_len", dtype=np.int32), ] mlmodel = ct.convert( traced_model, inputs=inputs, outputs=outputs, compute_units=ct.ComputeUnit.CPU_ONLY, minimum_deployment_target=ct.target.macOS14, ) save_path = output_path / "streaming_encoder.mlpackage" mlmodel.save(str(save_path)) print(f"Saved: {save_path}") # Save metadata metadata = { "model_id": args.model_id, "chunk_ms": args.chunk_ms, "mel_frames": chunk_size_in, "pre_cache_size": pre_cache_size, "valid_out_len": config["valid_out_len"], "samples_per_chunk": config["samples"], "hidden_dim": hidden_dim, "num_layers": num_layers, "cache_channel_size": cache_channel_size, "cache_time_size": cache_time_size, } metadata_path = output_path / "streaming_encoder_metadata.json" with open(metadata_path, "w") as f: json.dump(metadata, f, indent=2) print(f"Saved metadata: {metadata_path}") print(f"\n=== Export complete for {args.chunk_ms}ms ===") print(f"Output: {output_path}") print("\nNote: Decoder and Joint models are shared between chunk sizes.") print("Copy decoder.mlmodelc, joint_decision.mlmodelc, and vocab.json from 160ms directory.") if __name__ == "__main__": main()