Deploy Streaming Speaker Diarization with Sortformer2.1 on Modal

Published • guides
GPU

Adapted from Modal Examples’ sortformer2_1_speaker_diarization.py, available at https://github.com/modal-labs/modal-examples, under the MIT licence.

Speaker diarization answers “who spoke when?” in an audio stream. Unlike batch transcription, which processes a complete file after the fact, streaming diarization must emit speaker labels within a few hundred milliseconds of audio arriving. This guide walks through deploying NVIDIA’s Sortformer2.1 on Modal, a serverless GPU platform, to serve streaming speaker probabilities over WebSockets. Engineers building meeting summarizers, call analytics, or voice agents need this capability when transcripts alone are not enough: a speaker-aware transcript requires knowing which of four speakers is active at any given 80‑millisecond frame.

Modal’s pay‑per‑use billing makes this approach attractive for teams that cannot justify a dedicated GPU node. The service uses an L4 GPU, which is inexpensive per hour and has enough memory to hold Sortformer2.1 and run inference. The operational costs to watch are cold starts while the model downloads or loads, and data transfer if your audio source is not co‑located. For a broader look at voice stacks, see our coverage of Mistral’s Voxtral 3 voice agent SDK.

Prerequisites

  • Modal account and a configured local CLI. Create a Modal secret named huggingface-secret that contains a Hugging Face token, because Sortformer2.1 is gated.
  • Companion module sortformer2_1.py. The code below imports DiarizationConfig and NeMoStreamingDiarizer from that relative path, so copy the file from the Modal Examples repository into your project.
  • Dependencies and versions derived from the image definition: Python 3.12, CUDA 13.0.1, numpy<2, hf_transfer==0.1.9, huggingface_hub[hf-xet]==0.31.2, cuda-python==13.0.1, fastapi, and nemo_toolkit[asr] installed from the NeMo main git branch. System packages are git, libsndfile1, and ffmpeg.
  • GPU: the source uses an L4. You can change the gpu="L4" parameter, but keep your driver compatible with CUDA 13.0.1 if you use the same base image.

Step 1: Set up the Modal project and container image

The first block defines the app, a shared cache volume for Hugging Face models, and the runtime image. The base image is NVIDIA’s CUDA 13.0.1 development image with Python 3.12 added. The environment variables point Hugging Face and Torch caches to /model, which is mounted from the Modal volume, so the model does not have to be downloaded on every cold start. hf_transfer speeds up the initial download.

from pathlib import Path
from typing import Literal

import modal

app = modal.App("sortformer2-1-speaker-diarization")

CACHE_PATH = "/model"
cache_vol = modal.Volume.from_name("sortformer2_1-cache", create_if_missing=True)

image = (
    modal.Image.from_registry(
        "nvidia/cuda:13.0.1-cudnn-devel-ubuntu22.04", add_python="3.12"
    )
    .env(
        {
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
            "HF_HOME": CACHE_PATH,  # cache directory for Hugging Face models
            "CXX": "g++",
            "CC": "g++",
            "TORCH_HOME": CACHE_PATH,
        }
    )
    .apt_install("git", "libsndfile1", "ffmpeg")
    .uv_pip_install(
        "hf_transfer==0.1.9",
        "huggingface_hub[hf-xet]==0.31.2",
        "cuda-python==13.0.1",
        "numpy<2",
        "fastapi",
        "nemo_toolkit[asr]@git+https://github.com/NVIDIA/NeMo.git@main",
    )
)

with image.imports():
    import asyncio
    import json
    import time

    from fastapi import FastAPI, WebSocket, WebSocketDisconnect
    from starlette.websockets import WebSocketState

    from .sortformer2_1 import DiarizationConfig, NeMoStreamingDiarizer

The with image.imports() block pre‑declares imports so Modal can build the image without running the application. The source pins numpy<2; keep that constraint unless you test against a pinned NeMo commit, because changing it can produce a broken image.

Step 2: Define the diarization service class and configurations

This @app.cls block loads the model once in @modal.enter(), sets two latency profiles, and creates a FastAPI app inside the class. The low‑latency config is for real‑time use with small chunks and aggressive speaker cache refresh; the high‑latency config uses much longer context for higher accuracy at the cost of added delay.

@app.cls(
    image=image,
    volumes={CACHE_PATH: cache_vol},
    gpu="L4",
    secrets=[modal.Secret.from_name("huggingface-secret")],
)
class Sortformer2_1_Speaker_Diarization:
    @modal.enter()
    def enter(self):
        self._LOW_LATENCY_CONFIG = DiarizationConfig(
            max_num_speakers=4,
            chunk_len=6,
            chunk_right_context=7,
            fifo_len=188,
            spkcache_refresh_rate=144,
            spkcache_len=188,
        )
        self._HIGH_LATENCY_CONFIG = DiarizationConfig(
            max_num_speakers=4,
            chunk_len=340,
            chunk_right_context=40,
            fifo_len=40,
            spkcache_refresh_rate=300,
            spkcache_len=188,
        )
        self.latency: Literal["low", "high"] = "low"
        self._SORTFORMER_FRAME_SIZE_BYTES = (
            16000 * 0.08 * 2
        )  # sample rate * frame size in seconds * 2 bytes (16 bit)
        if self.latency == "low":
            self.config = self._LOW_LATENCY_CONFIG
        else:
            self.config = self._HIGH_LATENCY_CONFIG
        # load model from Hugging Face model card directly (You need a Hugging Face token)
        self.diarizer = NeMoStreamingDiarizer(
            cfg=self.config, model="nvidia/diar_streaming_sortformer_4spk-v2.1"
        )

        self.web_app = FastAPI()

        @self.web_app.websocket("/ws")
        async def run_with_websocket(ws: WebSocket):
            audio_queue = asyncio.Queue()
            output_queue = asyncio.Queue()

            async def recv_loop(ws, audio_queue):
                audio_buffer = bytearray()
                while True:
                    data = await ws.receive_bytes()
                    audio_buffer.extend(data)
                    if len(audio_buffer) > self._SORTFORMER_FRAME_SIZE_BYTES:
                        await audio_queue.put(audio_buffer)
                        audio_buffer = bytearray()

            async def inference_loop(audio_queue, output_queue):
                while True:
                    audio_data = await audio_queue.get()

                    start_time = time.perf_counter()
                    diar_result = self.diarizer.diarize(audio_data)

                    probs = self._get_speaker_probabilities(diar_result)
                    await output_queue.put(json.dumps(probs))

                    end_time = time.perf_counter()
                    print(
                        f"time taken to diarize audio segment: {end_time - start_time} seconds"
                    )

            async def send_loop(output_queue, ws):
                while True:
                    output = await output_queue.get()
                    print(f"sending diarize result: {output}")
                    await ws.send_text(output)

            await ws.accept()

            try:
                tasks = [
                    asyncio.create_task(recv_loop(ws, audio_queue)),
                    asyncio.create_task(inference_loop(audio_queue, output_queue)),
                    asyncio.create_task(send_loop(output_queue, ws)),
                ]
                await asyncio.gather(*tasks)
            except WebSocketDisconnect:
                print("WebSocket disconnected")
                ws = None
            except Exception as e:
                print("Exception:", e)
            finally:
                self.diarizer.reset_state()
                if ws and ws.application_state is WebSocketState.CONNECTED:
                    await ws.close(code=1011)  # internal error
                    ws = None
                for task in tasks:
                    if not task.done():
                        try:
                            task.cancel()
                            await task
                        except asyncio.CancelledError:
                            pass

The two configurations differ mainly in how much right context the model can use and how often the speaker cache is refreshed. Only one profile is active at a time, controlled by self.latency. The frame size calculation assumes 16 kHz mono 16‑bit PCM audio; if your client sends a different sample rate or stereo, the buffering logic will misalign.

ParameterLow latencyHigh latency
max_num_speakers44
chunk_len6 (0.48 s)340 (27.2 s)
chunk_right_context7 (0.56 s)40 (3.2 s)
fifo_len18840
spkcache_refresh_rate144300
spkcache_len188188

To switch profiles, change that string to "high" before deploying.

Step 3: Expose the inference endpoint and convert probabilities

The two methods after enter are small but essential. webapp is the Modal ASGI endpoint that returns the internal FastAPI app, and _get_speaker_probabilities reduces the model’s per‑frame speaker probability matrix to a single vector for the whole segment.

    @modal.asgi_app()
    def webapp(self):
        return self.web_app

    def _get_speaker_probabilities(self, spk_pred):
        # spk_pred is a 6x4 matrix of probabilities
        # We want to return a 1x4 vector of probabilities for the total time window
        # We can take the mean across the time dimension (axis 0)
        return spk_pred.mean(axis=0).tolist()

The mean across the time axis discards temporal detail. That is fine for a client that only needs the prominent speaker over a buffered segment, but it is not sufficient for a full speaker‑attributed transcript. To build that, you would need to send per‑frame probabilities from the model, not the averaged vector.

Step 4: Serve the frontend and wire the WebSocket URL

The second class serves a static frontend and injects the backend WebSocket URL into the page. It uses a smaller Debian‑slim image because no GPU is needed for serving HTML, JavaScript, and static files.

web_image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("fastapi")
    .add_local_dir(
        Path(__file__).parent / "streaming-diarization-frontend", "/root/frontend"
    )
)

with web_image.imports():
    from fastapi import FastAPI, WebSocket
    from fastapi.responses import HTMLResponse, Response
    from fastapi.staticfiles import StaticFiles


@app.cls(image=web_image)
@modal.concurrent(max_inputs=20)
class WebServer:
    @modal.asgi_app()
    def web(self):
        web_app = FastAPI()
        web_app.mount("/static", StaticFiles(directory="frontend"))

        @web_app.get("/status")
        async def status():
            return Response(status_code=200)

        # serve frontend
        @web_app.get("/")
        async def index():
            html_content = open("frontend/index.html").read()

            # Get the base WebSocket URL (without transcriber parameters)
            cls_instance = modal.Cls.from_name(
                "sortformer2-1-speaker-diarization", "Sortformer2_1_Speaker_Diarization"
            )()
            ws_base_url = (
                cls_instance.webapp.get_web_url().replace("http", "ws") + "/ws"
            )
            script_tag = f'<script>window.WS_BASE_URL = "{ws_base_url}";</script>'
            html_content = html_content.replace(
                '<script src="/static/sortformer2_1.js"></script>',
                f'{script_tag}\n<script src="/static/sortformer2_1.js"></script>',
            )
            return HTMLResponse(content=html_content)

        return web_app

The modal.concurrent(max_inputs=20) limit prevents a single frontend container from being overwhelmed by static‑file requests. When the index route is called, it dynamically builds a WebSocket URL from the backend’s webapp endpoint and injects it as a global variable into the frontend JavaScript.

What to watch out for

Shared state across connections. This backend class loads one NeMoStreamingDiarizer instance that lives for the life of the container. If two clients open WebSocket connections simultaneously, both will call self.diarizer.diarize() on the same object and both will call self.diarizer.reset_state() in their finally blocks. One client can wipe the speaker memory of another. The source does not guard against this. If you expect concurrency, run one container per connection or add a lock and per‑session diarizer instances.

Unbounded queues. audio_queue and output_queue have no maximum size. If inference falls behind real‑time, especially after switching to the high‑latency config, the queues will grow until the container runs out of memory. The low‑latency config is more forgiving, but you should monitor the timestamps printed by inference_loop and consider limiting queue size.

Bleeding‑edge dependencies. The image installs NeMo from the main git branch with no commit hash. This can break at any time as NVIDIA pushes changes. For production, fork or pin a specific commit, and test that numpy<2 still holds.

Audio format assumptions. The buffer splits incoming bytes whenever the buffer length exceeds 16000 * 0.08 * 2 = 2560 bytes, which is one 80 ms frame of 16 kHz mono 16‑bit PCM. If the client sends 8 kHz audio, resampled audio with a different format, or stereo interleaved data, the model will receive misaligned frames and produce invalid speaker labels.

Cold start and model cache. The first request after a cold start will wait for the model to be downloaded from Hugging Face into the volume if the cache is empty. Later starts can reuse the persisted model, but GPU initialization and modal.enter() model load still add latency. Use this setup for interactive or low‑traffic services until you have measured cold starts in your chosen region.

Where to go next

This guide gives you a working streaming diarization endpoint, but a complete speech pipeline still needs speech‑to‑text and speaker‑attributed transcript assembly. You can pair Sortformer2.1 with NeMo’s streaming ASR models and emit per‑frame diarization probabilities to a client that aligns words with speakers. Modal’s WebSocket support also lets you deploy the same pattern with autoscaling and GPU snapshots to cut cold starts. For more on the voice agent layer that consumes these signals, see our piece on Mistral’s Voxtral 3 SDK.

Frequently asked questions

What is speaker diarization?

Speaker diarization identifies who spoke when in an audio stream. Unlike transcription, it does not convert speech to text; it labels speaker turns or frames. Sortformer2.1 outputs per-frame speaker probabilities for up to four speakers.

How do I deploy Sortformer2.1 on Modal?

Clone the Modal Examples repository and copy the companion `sortformer2_1.py` module into your project. Create a Modal secret named `huggingface-secret` with a gated Hugging Face token, then run `modal deploy`. The example builds a CUDA 13.0.1 image with NeMo and runs inference on an L4 GPU.

What audio format does this diarization example expect?

The buffering logic assumes 16 kHz mono 16-bit PCM audio. Each 80 ms frame is 2560 bytes. If you send another sample rate, stereo interleaved audio, or encoded audio, the frames will be misaligned and the model can produce invalid speaker labels.

How do I switch between low-latency and high-latency diarization?

Change `self.latency` from `"low"` to `"high"` inside the `enter` method before deploying. The high-latency profile uses a 340-frame chunk and 40-frame right context for more accuracy, while the low-latency profile uses a 6-frame chunk and 7-frame right context.

Why is the first request slow after a Modal cold start?

The first container must initialize the L4 GPU, load the diarization model, and download it from Hugging Face if the `/model` volume is empty. Later cold starts can reuse the cached model, but GPU and model initialization still add latency. Measure your own region and image before relying on interactive response times.

Can the example handle multiple WebSocket clients at once?

The provided backend uses one shared `NeMoStreamingDiarizer` instance, so concurrent connections can call `diarize` and `reset_state` on the same object. One client can wipe another client's speaker memory. For concurrency, run one container per connection or add locking and per-session diarizer state.

Free interactive tools for the decisions this piece raises.

Related Guides