Deploy StreamDiffusionV2 Video Style Transfer with WMA

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

class Prompt(BaseModel):
    type: Literal["prompt"] = "prompt"
    text: str = Field(
        min_length=1,
        max_length=400,
        description="Style prompt. The stream restarts to apply it.",
    )
# ...
ServerMessage = Annotated[Union[Status, Rejected], Field(discriminator="type")]

class RestyleApp(fal.wma.App, python_version="3.11"):
    machine_type = "GPU-H100"
    requirements = [
        "fal[wma]==1.82.1",
# ...
        "aiortc",
    ]
    startup_timeout = 1800  # Weights download and kernel compilation take minutes.
    min_concurrency = 0
    max_concurrency = 2
    max_multiplexing = 1  # One session owns the GPU and the model state.
    keep_alive = 300

    realtime_contract = fal.wma.RealtimeContract(
        media=fal.wma.MediaContract(
            send=(
                fal.wma.Track(
                    kind="video",
                    source="camera",
                    required=True,
                    width=fal.wma.Constraint(ideal=640),
                    height=fal.wma.Constraint(ideal=480),
                    frame_rate=fal.wma.Constraint(ideal=30),
                ),
            ),
            receive=(
                fal.wma.Track(kind="video", width=SIZE, height=SIZE, frame_rate=FPS),
            ),
        ),
        client_messages=Prompt,
        server_messages=ServerMessage,

    def on_prompt(self, message):
        try:
            prompt = Prompt.model_validate(message)
        except ValueError:
            self.session.send(
                Rejected(reason="text must be 1 to 400 characters").model_dump()
            )
            return
        with self.cond:
            self.prompt = prompt.text
            self.cond.notify_all()

        def connect(pc):
            attached = False

            @pc.on("track")
            def incoming(track):
                nonlocal attached
                if track.kind == "video" and not attached:
                    attached = True
                    pc.addTrack(
                        fal.wma.make_video_queue_track(
                            restyle.out, fps=FPS, rebase_after_stall=True
                        )
                    )
                    restyle.start(track)

        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,
        )

fal run fal_demos/video/wma/video_restyle.py::RestyleApp

fal deploy fal_demos/video/wma/video_restyle.py::RestyleApp --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 connectRestyle(
  appId: string,
  output: HTMLVideoElement,
  onEnded: (reason: string) => void,
): Promise<{ setPrompt(text: string): void; close(): Promise<void> }> {
  const camera = await navigator.mediaDevices.getUserMedia({
    video: true,
    audio: false,
  });
  const stopCamera = () => camera.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;
    stopCamera();
    output.srcObject = null;
    onEnded(reason);
  };
  const session = fal.realtime.open(wma(appId), {
    localStream: camera,
    receive: ["video"],
    onMedia(stream) {
      output.srcObject = stream;
      void output.play().catch(console.error);
    },
    onData: (raw) => {
      const msg = JSON.parse(raw);
      switch (msg.type) {
        case "status":
          console.log(msg.state, msg.output_frames, 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) {
    stopCamera();
    await session.close();
    throw error;
  }
  return {
    setPrompt: (text) => session.send({ type: "prompt", text }),
    async close() {
      try {
        await session.close();
      } finally {
        stopCamera();
        output.srcObject = null;
      }
    },
  };
}

    async def read_camera(self, track):
        # Drain the camera track so its queue never grows. Keep only the newest frame.
        while True:
            try:
                self.latest = await track.recv()
            except Exception:
                break
        # The camera ended. Stop feeding frames so the model stops generating.
        self.latest = None
        with self.cond:
            self.frames.clear()
        self.set_state("waiting_for_camera")

    def put_output(self, frame):
        """Queue a frame for playback on the event loop. Count the frame it displaces."""
        if self.out.full():
            self.dropped_output_frames += 1
        fal.wma.queue_put_drop_oldest(self.out, frame)