Deploy Moshi Real-time Voice Chat with WMA

git clone https://github.com/fal-ai-community/fal-demos.git
cd fal-demos
uv venv --python 3.12
source .venv/bin/activate
uv pip install 'fal[wma]==1.82.1'
fal auth login

class Reset(BaseModel):
    type: Literal["reset"] = "reset"


class Text(BaseModel):
    type: Literal["text"] = "text"
    text: str = Field(description="One piece of what Moshi is saying, in order.")
# ...
ServerMessage = Annotated[Union[Text, Status, Rejected], Field(discriminator="type")]

class VoiceChatApp(fal.wma.App, python_version="3.12"):
    machine_type = "GPU-H100"
    requirements = [
        "fal[wma]==1.82.1",
        "moshi==0.2.13",
# ...
        "aiortc",
    ]
    startup_timeout = 1200  # Weights download and CUDA graph capture take minutes.
    min_concurrency = 0
    max_concurrency = 2
    max_multiplexing = 1  # One session owns the GPU and the streaming state.
    keep_alive = 300

    realtime_contract = fal.wma.RealtimeContract(
        media=fal.wma.MediaContract(
            send=(fal.wma.Track(kind="audio", source="microphone", required=True),),
            receive=(fal.wma.Track(kind="audio"),),
        ),
        client_messages=Reset,
        server_messages=ServerMessage,
        client_message_examples=(
            fal.wma.MessageExample(name="New conversation", payload={"type": "reset"}),
        ),
    )

    def step(self, chunk, skip=False):
        """Feed one input frame. Return (pcm, text) pairs for the frames Moshi produced.

        `chunk` is float32 audio shaped (1, 1, FRAME_SAMPLES). With `skip`, the
        frame is encoded and then discarded, as upstream does for the first frame.
        """
        codes = self.mimi.encode(chunk)
        if skip:
            # The first encoded frame carries the encoder's left padding. Upstream
            # resets the encoder so the next frame gets the padding again.
            self.mimi.reset_streaming()
        out = []
        for c in range(codes.shape[-1]):
            tokens = self.lm_gen.step(codes[:, :, c : c + 1])
            if tokens is None:  # Moshi has not produced a frame yet.
                continue
            pcm = self.mimi.decode(tokens[:, 1:])[0, 0].float().cpu().numpy()
            token = tokens[0, 0, 0].item()
            text = None
            if token not in TEXT_PADDING_TOKENS:
                text = self.tokenizer.id_to_piece(token).replace("▁", " ")
            out.append((pcm, text))
        return out

        def connect(pc):
            attached = False

            @pc.on("track")
            def incoming(track):
                nonlocal attached
                if track.kind == "audio" and not attached:
                    attached = True
                    speaker = make_audio_queue_track(voice.out)
                    pc.addTrack(speaker)
                    voice.start(track, speaker)

        async def expire():
            await asyncio.sleep(MAX_SESSION_SECONDS)
            await session.close()

        session.create_task(expire())
        return fal.wma.AiortcPeer(
            session,
            connect,
            rtc_configuration=RTCConfiguration(iceServers=servers),
            disconnected_grace_seconds=10,
        )

        async def recv(self):
            now = time.monotonic()
            if self.start is None:
                self.start = now
            due = self.start + self.pts / SAMPLE_RATE
            if due < now - 1:  # The consumer stalled. Do not burst to catch up.
                self.start = now - self.pts / SAMPLE_RATE
            elif due > now:
                await asyncio.sleep(due - now)
            while len(self.pending) < PACKET_SAMPLES and not queue.empty():
                self.pending = np.concatenate([self.pending, queue.get_nowait()])
            if len(self.pending) >= PACKET_SAMPLES:
                chunk = self.pending[:PACKET_SAMPLES]
                self.pending = self.pending[PACKET_SAMPLES:]
            else:
                chunk = np.zeros(PACKET_SAMPLES, dtype=np.float32)

fal run fal_demos/video/wma/voice_chat.py::VoiceChatApp

fal deploy fal_demos/video/wma/voice_chat.py::VoiceChatApp --auth private

npm install @fal-ai/client@1.11.0-alpha.5 @fal-ai/server-proxy@1.2.1

import { fal } from "@fal-ai/client";
import { wma } from "@fal-ai/client/realtime/wma";

fal.config({ proxyUrl: "/api/fal/proxy" });

export async function connectVoiceChat(
  appId: string,
  speaker: HTMLAudioElement,
  onText: (text: string) => void,
  onEnded: (reason: string) => void,
): Promise<{ reset(): void; close(): Promise<void> }> {
  const microphone = await navigator.mediaDevices.getUserMedia({
    audio: { echoCancellation: true, noiseSuppression: true },
    video: false,
  });
  const stopMicrophone = () => microphone.getTracks().forEach((track) => track.stop());
  let ended = false;
  const finish = (reason: string) => {
    // A failure reaches both onError and onState, so run the cleanup once.
    if (ended) return;
    ended = true;
    stopMicrophone();
    speaker.srcObject = null;
    onEnded(reason);
  };
  const session = fal.realtime.open(wma(appId), {
    localStream: microphone,
    receive: ["audio"],
    onMedia(stream) {
      speaker.srcObject = stream;
      void speaker.play().catch(console.error);
    },
    onData: (raw) => {
      const msg = JSON.parse(raw);
      switch (msg.type) {
        case "text":
          onText(msg.text);
          break;
        case "status":
          console.log(msg.state, msg.heard_seconds, msg.dropped_input_frames);
          break;
        case "rejected":
          console.error(msg.reason);
          break;
      }
    },
    onState(state) {
      if (state === "closed" || state === "failed") finish(state);
    },
    onError(error) {
      console.error(error);
      finish("failed");
    },
  });
  try {
    await session.ready;
  } catch (error) {
    stopMicrophone();
    await session.close();
    throw error;
  }
  return {
    reset: () => session.send({ type: "reset" }),
    async close() {
      try {
        await session.close();
      } finally {
        stopMicrophone();
        speaker.srcObject = null;
      }
    },
  };
}

    def put_output(self, pcm):
        """Queue audio for playback on the event loop. Count the audio it displaces."""
        if self.out.full():
            self.dropped_output_samples += len(self.out._queue[0])
        fal.wma.queue_put_drop_oldest(self.out, pcm)