Function main

examples/models/sortformer/export_sortformer.py:356–405

()

def main():
    parser = argparse.ArgumentParser(
        description="Export streaming Sortformer diarizer to ExecuTorch"
    )
    parser.add_argument("--output-dir", default="./sortformer_exports")
    parser.add_argument(
        "--nemo-path",
        type=str,
        help="Path to .nemo model file",
    )
    parser.add_argument(
        "--hf-model",
        type=str,
        help="HuggingFace model ID (default: nvidia/diar_streaming_sortformer_4spk-v2)",
    )
    parser.add_argument(
        "--backend",
        type=str,
        default="xnnpack",
        choices=["portable", "xnnpack", "cuda", "cuda-windows"],
        help="Backend for acceleration (default: xnnpack)",
    )

    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    print("Loading model...")
    model = load_model(nemo_path=args.nemo_path, hf_model=args.hf_model)

    print("\nExporting components...")
    programs, metadata = export_all(model, backend=args.backend)

    et = lower_to_executorch(programs, metadata=metadata, backend=args.backend)

    pte_path = os.path.join(args.output_dir, "sortformer.pte")
    print(f"\nSaving ExecuTorch program to: {pte_path}")
    with open(pte_path, "wb") as f:
        et.write_to_file(f)

    if args.backend in ("cuda", "cuda-windows"):
        cuda_blob_path = os.path.join(args.output_dir, "aoti_cuda_blob.ptd")
        print(f"Writing CUDA named-data blob to: {cuda_blob_path}")
        et.write_tensor_data_to_file(args.output_dir)
        if not os.path.isfile(cuda_blob_path):
            raise FileNotFoundError(
                f"Expected CUDA named-data blob at {cuda_blob_path}, but it was not created."
            )
    print(f"Saved {os.path.getsize(pte_path) / (1024 * 1024):.1f} MB")

    print("\nDone!")


if __name__ == "__main__":