diff --git a/flashdreams/flashdreams/api_v2/client_window.py b/flashdreams/flashdreams/api_v2/client_window.py index 62b829245..8befe1c02 100644 --- a/flashdreams/flashdreams/api_v2/client_window.py +++ b/flashdreams/flashdreams/api_v2/client_window.py @@ -20,7 +20,7 @@ class IClientWindow(InputSource, OutputSink, ABC): is given that description in :meth:`OutputSink.open`. One thread makes every call on a window, so an implementation needs no - locking of its own. + locking except when its backend delivers input from another thread. Created by the runtime, never by an application. """ diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py new file mode 100644 index 000000000..a985ee377 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Application lifecycle runner for the v2 runtime.""" + +from collections.abc import Sequence + +from flashdreams.api_v2.application import IApplication +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_runner import run_session + + +class ApplicationRunner: + """Create and run one application session against one client window.""" + + def __init__(self, application: IApplication, client_window: IClientWindow) -> None: + """ + Args: + application: Long-lived application that creates the session. + client_window: Window that supplies input and presents generated output. + """ + self._application = application + self._client_window = client_window + + def run( + self, + session_desc: SessionDesc, + commandline_args: Sequence[str] = (), + ) -> None: + """Initialize the application, create one session, and run it. + + The application is closed before this method returns or raises. + + Args: + session_desc: Output shape and timing requested for the session. + commandline_args: Arguments owned and parsed by the application. + """ + try: + self._application.init(commandline_args) + session = self._application.create_session(session_desc) + run_session(session, self._client_window) + finally: + self._application.close() diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py new file mode 100644 index 000000000..c5f2327bb --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Create v2 client windows from runtime arguments.""" + +import argparse + +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow + + +def create_client_window(parsed_args: argparse.Namespace) -> IClientWindow: + """Create the client window selected by the presentation mode. + + Args: + parsed_args: Runtime arguments. Mode-specific fields are read only by + the selected mode. + + Returns: + Client window for the selected mode. + + Raises: + ValueError: ``mode`` is unsupported. + """ + if parsed_args.mode == "webrtc": + return WebRTCClientWindow(host=parsed_args.host, port=parsed_args.port) + raise ValueError(f"Unsupported client-window mode: {parsed_args.mode!r}.") diff --git a/flashdreams/flashdreams/runtime_v2/serving/__init__.py b/flashdreams/flashdreams/runtime_v2/serving/__init__.py new file mode 100644 index 000000000..3b470c930 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/serving/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serving backends for the v2 runtime.""" diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js new file mode 100644 index 000000000..029e714e7 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js @@ -0,0 +1,61 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +const peer = new RTCPeerConnection(); +const controls = peer.createDataChannel("controls"); +peer.addTransceiver("video", {direction: "recvonly"}); + +peer.ontrack = event => { + document.getElementById("video").srcObject = + event.streams[0] ?? new MediaStream([event.track]); +}; + +const send = payload => { + if (controls.readyState === "open") { + controls.send(JSON.stringify(payload)); + } +}; + +window.addEventListener("keydown", event => { + send({type: "keyboard", key: event.key, pressed: true}); +}); + +window.addEventListener("keyup", event => { + send({type: "keyboard", key: event.key, pressed: false}); +}); + +let activationPressed = false; +document.getElementById("activate").onclick = event => { + activationPressed = !activationPressed; + event.currentTarget.textContent = + activationPressed ? "Deactivate" : "Activate"; + send({type: "keyboard", key: "r", pressed: activationPressed}); +}; + +document.getElementById("reset").onclick = () => { + send({type: "reset"}); +}; + +window.addEventListener("beforeunload", () => send({type: "close"})); + +async function connect() { + while (true) { + const health = await fetch("/healthz"); + if (health.ok && (await health.json()).open) { + break; + } + await new Promise(resolve => setTimeout(resolve, 100)); + } + await peer.setLocalDescription(await peer.createOffer()); + const response = await fetch("/api/webrtc/offer", { + method: "POST", + headers: {"content-type": "application/json"}, + body: JSON.stringify(peer.localDescription), + }); + if (!response.ok) { + throw new Error(await response.text()); + } + await peer.setRemoteDescription(await response.json()); +} + +connect(); diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/index.html b/flashdreams/flashdreams/runtime_v2/serving/web/index.html new file mode 100644 index 000000000..2fb07c96e --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/serving/web/index.html @@ -0,0 +1,18 @@ + + + + + + + FlashDreams WebRTC + + + + + + + + diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py new file mode 100644 index 000000000..68041879d --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -0,0 +1,492 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone WebRTC server used by the v2 client window.""" + +from __future__ import annotations + +import asyncio +import json +import socket +import threading +import time +from collections.abc import Callable +from fractions import Fraction +from importlib.resources import files +from typing import Any + +import numpy as np +import torch +from aiohttp import web +from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription +from aiortc.mediastreams import MediaStreamError +from av import VideoFrame + +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + CloseUserInputEventData, + KeyboardUserInputEventData, + ResetUserInputEventData, + UserInputEvent, +) +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +_WEB_RESOURCES = files("flashdreams.runtime_v2.serving").joinpath("web") +_BROWSER_PAGE = _WEB_RESOURCES.joinpath("index.html").read_text(encoding="utf-8") +_BROWSER_SCRIPT = _WEB_RESOURCES.joinpath("app.js").read_text(encoding="utf-8") + + +class _VideoTrack(MediaStreamTrack): + """Video track whose frames are supplied by the server.""" + + kind = "video" + + def __init__(self, frames_per_second: int) -> None: + super().__init__() + self._frames_per_second = frames_per_second + self._time_base = Fraction(1, frames_per_second) + self._frames: asyncio.Queue[np.ndarray[Any, np.dtype[np.uint8]] | None] = ( + asyncio.Queue() + ) + self._next_frame_time: float | None = None + self._pts = 0 + self._closed = False + + async def enqueue( + self, frames: tuple[np.ndarray[Any, np.dtype[np.uint8]], ...] + ) -> None: + """Append generated RGB frames for the WebRTC sender.""" + if self._closed: + return + for frame in frames: + await self._frames.put(frame) + + async def recv(self) -> VideoFrame: + """Return the next generated frame when aiortc requests one.""" + if self._closed: + raise MediaStreamError + frame = await self._frames.get() + if frame is None: + raise MediaStreamError + + loop = asyncio.get_running_loop() + now = loop.time() + if self._next_frame_time is None: + self._next_frame_time = now + else: + self._next_frame_time += 1.0 / self._frames_per_second + await asyncio.sleep(max(0.0, self._next_frame_time - now)) + + video_frame = VideoFrame.from_ndarray(frame, format="rgb24") + video_frame.pts = self._pts + video_frame.time_base = self._time_base + self._pts += 1 + return video_frame + + async def close(self) -> None: + """Stop the track and release a pending receiver.""" + if self._closed: + return + self._closed = True + self._frames.put_nowait(None) + self.stop() + + +class WebRTCServer: + """Own the HTTP, signaling, input buffering, and media transport.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 0, + startup_timeout_seconds: float = 10.0, + ) -> None: + """ + Args: + host: Interface on which the HTTP server listens. + port: Listening port. Zero asks the operating system to choose one. + startup_timeout_seconds: Maximum time to wait for server startup. + + Raises: + RuntimeError: The server cannot start. + TimeoutError: The server does not start before the timeout. + """ + if not host: + raise ValueError("host must not be empty.") + if port < 0 or port > 65535: + raise ValueError("port must be between 0 and 65535.") + if startup_timeout_seconds <= 0: + raise ValueError("startup_timeout_seconds must be > 0.") + + self._host = host + self._port = port + self._startup_timeout_seconds = startup_timeout_seconds + self._input_callback: Callable[[UserInputEvent], None] | None = None + self._started = threading.Event() + self._startup_error: BaseException | None = None + self._loop: asyncio.AbstractEventLoop | None = None + self._runner: web.AppRunner | None = None + self._peer_connection: RTCPeerConnection | None = None + self._video_track: _VideoTrack | None = None + self._session_desc: SessionDesc | None = None + self._session_start_ns: int | None = None + self._closed = False + self._client_connected = False + self._thread = threading.Thread( + target=self._run_server, + name="flashdreams-webrtc", + daemon=True, + ) + self._thread.start() + if not self._started.wait(startup_timeout_seconds): + raise TimeoutError("WebRTC server did not start before the timeout.") + if self._startup_error is not None: + raise RuntimeError( + "WebRTC server failed to start." + ) from self._startup_error + + @property + def host(self) -> str: + """Return the interface on which the server is listening.""" + return self._host + + @property + def port(self) -> int: + """Return the bound server port.""" + return self._port + + @property + def url(self) -> str: + """Return the browser URL for this server.""" + return f"http://{self._host}:{self._port}/" + + def open(self, session_desc: SessionDesc) -> None: + """Configure the server for one session's generated video. + + Args: + session_desc: Resolved dimensions, frame rate, and tensor layout. + + Raises: + RuntimeError: The server is closed or already open. + """ + if self._closed: + raise RuntimeError("Cannot open a closed WebRTC server.") + if self._session_desc is not None: + raise RuntimeError("WebRTC server is already open.") + if self._input_callback is None: + raise RuntimeError("Register an input callback before opening WebRTC.") + self._session_desc = session_desc + self._session_start_ns = time.monotonic_ns() + + def register_input_callback( + self, callback: Callable[[UserInputEvent], None] + ) -> None: + """Register the function called for each received browser event. + + Args: + callback: Function that accepts one validated, timestamped event. + + Raises: + RuntimeError: A callback has already been registered. + """ + if self._input_callback is not None: + raise RuntimeError("An input callback is already registered.") + self._input_callback = callback + + def write(self, result: StepResult) -> None: + """Deliver one generated result to the browser's video track. + + Args: + result: Generated frames matching the description passed to + :meth:`open`. + + Raises: + RuntimeError: The server is not open or has been closed. + ValueError: The result shape or layout does not match the session. + """ + if self._closed: + raise RuntimeError("Cannot write to a closed WebRTC server.") + session_desc = self._session_desc + if session_desc is None: + raise RuntimeError("Open the WebRTC server before writing.") + frames = _result_to_rgb_frames(result, session_desc) + loop = self._loop + if loop is None: + raise RuntimeError("WebRTC server is not running.") + future = asyncio.run_coroutine_threadsafe(self._enqueue_frames(frames), loop) + future.result() + + def close(self) -> None: + """Close the peer connection and stop the WebRTC server.""" + if self._closed: + return + self._closed = True + loop = self._loop + if loop is None: + return + future = asyncio.run_coroutine_threadsafe(self._shutdown(), loop) + future.result(timeout=self._startup_timeout_seconds) + loop.call_soon_threadsafe(loop.stop) + self._thread.join(timeout=self._startup_timeout_seconds) + if self._thread.is_alive(): + raise TimeoutError("WebRTC server did not stop before the timeout.") + + def _run_server(self) -> None: + """Own the WebRTC asyncio loop for the lifetime of the server.""" + loop = asyncio.new_event_loop() + self._loop = loop + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(self._start_server()) + except BaseException as error: + self._startup_error = error + self._started.set() + loop.close() + return + self._started.set() + try: + loop.run_forever() + finally: + loop.close() + + async def _start_server(self) -> None: + """Create and bind the standalone aiohttp application.""" + app = web.Application() + app.router.add_get("/", self._serve_browser) + app.router.add_get("/app.js", self._serve_browser_script) + app.router.add_get("/healthz", self._health) + app.router.add_post("/api/webrtc/offer", self._offer) + runner = web.AppRunner(app) + await runner.setup() + address_family = socket.AF_INET6 if ":" in self._host else socket.AF_INET + server_socket = socket.socket(address_family, socket.SOCK_STREAM) + try: + server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server_socket.bind((self._host, self._port)) + server_socket.setblocking(False) + server_socket.listen(128) + self._port = int(server_socket.getsockname()[1]) + site = web.SockSite(runner, server_socket) + await site.start() + except Exception: + server_socket.close() + await runner.cleanup() + raise + self._runner = runner + + async def _serve_browser(self, _: web.Request) -> web.Response: + """Return the minimal browser client.""" + return web.Response(text=_BROWSER_PAGE, content_type="text/html") + + async def _serve_browser_script(self, _: web.Request) -> web.Response: + """Return the browser client's JavaScript.""" + return web.Response(text=_BROWSER_SCRIPT, content_type="text/javascript") + + async def _health(self, _: web.Request) -> web.Response: + """Report whether the server has an open session and client.""" + return web.json_response( + { + "open": self._session_desc is not None, + "client_connected": self._client_connected, + } + ) + + async def _offer(self, request: web.Request) -> web.Response: + """Negotiate one browser peer connection.""" + if self._closed: + raise web.HTTPServiceUnavailable(reason="WebRTC server is closed.") + session_desc = self._session_desc + if session_desc is None: + raise web.HTTPConflict(reason="WebRTC server is not open.") + if self._peer_connection is not None: + raise web.HTTPConflict(reason="A WebRTC client is already connected.") + + try: + payload = await request.json() + except (json.JSONDecodeError, web.HTTPException) as error: + raise web.HTTPBadRequest(reason="Expected a JSON WebRTC offer.") from error + if not isinstance(payload, dict): + raise web.HTTPBadRequest(reason="WebRTC offer must be an object.") + sdp = payload.get("sdp") + offer_type = payload.get("type") + if not isinstance(sdp, str) or not isinstance(offer_type, str): + raise web.HTTPBadRequest( + reason="WebRTC offer requires string sdp and type." + ) + + peer_connection = RTCPeerConnection() + video_track = _VideoTrack(session_desc.frames_per_second_for_ui) + peer_connection.addTrack(video_track) + self._peer_connection = peer_connection + self._video_track = video_track + + @peer_connection.on("datachannel") + def on_datachannel(channel: Any) -> None: + self._client_connected = True + + @channel.on("message") + def on_message(message: Any) -> None: + try: + self._buffer_browser_message(message) + except ValueError as error: + channel.send(json.dumps({"type": "error", "message": str(error)})) + + @channel.on("close") + def on_close() -> None: + self._record_client_disconnect() + + @peer_connection.on("connectionstatechange") + async def on_connectionstatechange() -> None: + if peer_connection.connectionState in {"failed", "disconnected", "closed"}: + self._record_client_disconnect() + + try: + await peer_connection.setRemoteDescription( + RTCSessionDescription(sdp=sdp, type=offer_type) + ) + await peer_connection.setLocalDescription( + await peer_connection.createAnswer() + ) + except Exception: + self._peer_connection = None + self._video_track = None + await video_track.close() + await peer_connection.close() + raise + + local_description = peer_connection.localDescription + if local_description is None: + raise web.HTTPInternalServerError( + reason="WebRTC peer did not create an answer." + ) + return web.json_response( + {"sdp": local_description.sdp, "type": local_description.type} + ) + + def _buffer_browser_message(self, raw_message: object) -> None: + """Validate and append one data-channel message.""" + if not isinstance(raw_message, str): + raise ValueError("Browser event must be a JSON string.") + try: + payload = json.loads(raw_message) + except json.JSONDecodeError as error: + raise ValueError("Browser event must contain valid JSON.") from error + if not isinstance(payload, dict): + raise ValueError("Browser event must be a JSON object.") + + event_type = payload.get("type") + if event_type == "keyboard": + key = payload.get("key") + pressed = payload.get("pressed") + if not isinstance(key, str) or not key: + raise ValueError("Keyboard event requires a non-empty key.") + if not isinstance(pressed, bool): + raise ValueError("Keyboard event requires a boolean pressed value.") + event_data = KeyboardUserInputEventData(key=key, pressed=pressed) + elif event_type == "reset": + event_data = ResetUserInputEventData() + elif event_type == "close": + event_data = CloseUserInputEventData() + else: + raise ValueError( + "Browser event type must be 'keyboard', 'reset', or 'close'." + ) + self._append_event(event_data) + + def _append_event( + self, + event_data: ( + KeyboardUserInputEventData + | ResetUserInputEventData + | CloseUserInputEventData + ), + ) -> None: + """Timestamp and buffer one validated browser event.""" + session_start_ns = self._session_start_ns + if session_start_ns is None: + return + timestamp_us = np.uint64((time.monotonic_ns() - session_start_ns) // 1_000) + event = UserInputEvent(timestamp=timestamp_us, event_data=event_data) + callback = self._input_callback + if callback is None: + raise RuntimeError("WebRTC input callback is not registered.") + # Pass that UserInputEvent to the callback. + # The callback stores it in WebRTCClientWindow’s thread-safe queue. + callback(event) + + def _record_client_disconnect(self) -> None: + """Buffer one close event when the active browser disconnects.""" + if not self._client_connected: + return + self._client_connected = False + if not self._closed: + self._append_event(CloseUserInputEventData()) + + async def _enqueue_frames( + self, frames: tuple[np.ndarray[Any, np.dtype[np.uint8]], ...] + ) -> None: + """Append frames to the active media track, if connected.""" + track = self._video_track + if track is not None: + await track.enqueue(frames) + + async def _shutdown(self) -> None: + """Release async server resources on their owning loop.""" + peer_connection = self._peer_connection + self._peer_connection = None + track = self._video_track + self._video_track = None + if track is not None: + await track.close() + if peer_connection is not None: + await peer_connection.close() + runner = self._runner + self._runner = None + if runner is not None: + await runner.cleanup() + + +def _result_to_rgb_frames( + result: StepResult, session_desc: SessionDesc +) -> tuple[np.ndarray[Any, np.dtype[np.uint8]], ...]: + """Convert one result to time-major RGB uint8 frames.""" + output = result.output.detach() + if result.output_layout == VideoTensorLayout.tchw: + frames = output + elif result.output_layout == VideoTensorLayout.btchw: + if output.ndim != 5 or output.shape[0] != 1: + raise ValueError("btchw WebRTC output requires a batch size of one.") + frames = output[0] + elif result.output_layout == VideoTensorLayout.bcthw: + if output.ndim != 5 or output.shape[0] != 1: + raise ValueError("bcthw WebRTC output requires a batch size of one.") + frames = output[0].permute(1, 0, 2, 3) + elif result.output_layout == VideoTensorLayout.bvtchw: + if output.ndim != 6 or output.shape[:2] != (1, 1): + raise ValueError( + "bvtchw WebRTC output requires one batch and one video view." + ) + frames = output[0, 0] + else: + raise ValueError(f"Unsupported WebRTC output layout: {result.output_layout}.") + + if frames.ndim != 4: + raise ValueError("WebRTC output must resolve to a tchw tensor.") + if frames.shape[0] != result.frame_count: + raise ValueError("StepResult.frame_count does not match its output tensor.") + if frames.shape[1] not in (1, 3): + raise ValueError("WebRTC output must have one or three color channels.") + if frames.shape[2:] != (session_desc.video_height, session_desc.video_width): + raise ValueError("WebRTC output dimensions do not match SessionDesc.") + if result.output_layout != session_desc.output_layout: + raise ValueError("StepResult.output_layout does not match SessionDesc.") + + if frames.shape[1] == 1: + frames = frames.repeat(1, 3, 1, 1) + if frames.is_floating_point(): + frames = ((frames.to(torch.float32).clamp(-1.0, 1.0) + 1.0) * 127.5).round() + frames = frames.clamp(0, 255).to(torch.uint8) + frames = frames.permute(0, 2, 3, 1).contiguous().cpu() + return tuple(np.asarray(frame.numpy()) for frame in frames) diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py new file mode 100644 index 000000000..28936d1fc --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""WebRTC client window for the v2 runtime.""" + +import queue + +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.serving.webrtc_server import WebRTCServer +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import UserInputEvent +from flashdreams.runtime_v2.user_input_events import UserInputEvents + + +class WebRTCClientWindow(IClientWindow): + """Implement ``IClientWindow`` with WebRTC input and presentation.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 0, + startup_timeout_seconds: float = 10.0, + ) -> None: + """Create the WebRTC backend. + + Construction is specific to this implementation; it is not part of the + ``IClientWindow`` protocol. + + Args: + host: Interface on which the HTTP server listens. + port: Listening port. Zero asks the operating system to choose one. + startup_timeout_seconds: Maximum time to wait for server startup. + """ + self._input_events: queue.SimpleQueue[UserInputEvent] = queue.SimpleQueue() + self.server = WebRTCServer( + host=host, + port=port, + startup_timeout_seconds=startup_timeout_seconds, + ) + + def handle_input(event: UserInputEvent) -> None: + """Buffer one backend event for the ``InputSource`` protocol.""" + self._input_events.put(event) + + self.server.register_input_callback(handle_input) + + def open(self, session_desc: SessionDesc) -> None: + """Implement ``OutputSink.open`` by configuring WebRTC output. + + Args: + session_desc: Resolved dimensions, frame rate, and tensor layout. + """ + self.server.open(session_desc) + + def get_user_input_events(self) -> UserInputEvents: + """Implement ``InputSource.get_user_input_events`` for browser input. + + Returns: + Buffered browser events in timestamp order, each returned once. + """ + events = [] + while True: + try: + events.append(self._input_events.get_nowait()) + except queue.Empty: + return UserInputEvents(events) + + def write(self, result: StepResult) -> None: + """Implement ``OutputSink.write`` by delivering a result to the browser. + + Args: + result: Generated frames matching the opened session. + """ + self.server.write(result) + + def close(self) -> None: + """Implement ``OutputSink.close`` by releasing WebRTC resources.""" + self.server.close() diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 5a0636e33..0206e6005 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -143,6 +143,7 @@ exclude = ["tests", "flashdreams._pytest_plugins*"] [tool.setuptools.package-data] "flashdreams.serving.webrtc" = ["web/*.html", "web/*.css", "web/*.js", "web/assets/*.svg"] +"flashdreams.runtime_v2.serving" = ["web/*.html", "web/*.js"] [dependency-groups] # Default CUDA 13 profile. No source binding -- Linux falls through to diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py new file mode 100644 index 000000000..0a26f8823 --- /dev/null +++ b/flashdreams/test_v2/test_application_runner.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the v2 application runner.""" + +from collections.abc import Sequence + +import pytest +import torch +from numpy import uint64 + +from flashdreams.api_v2.application import IApplication +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2.application_runner import ApplicationRunner +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + CloseUserInputEventData, + UserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +pytestmark = pytest.mark.ci_cpu + + +class _Session(ISession): + def __init__(self, session_desc: SessionDesc, calls: list[str]) -> None: + self._session_desc = session_desc + self._calls = calls + + def init(self) -> None: + self._calls.append("session.init") + + @property + def session_desc(self) -> SessionDesc: + return self._session_desc + + def step(self, step_index: int, events: UserInputEvents) -> StepResult: + del events + self._calls.append(f"session.step({step_index})") + return StepResult( + step_index=step_index, + output=torch.zeros((1, 3, 1, 2, 2)), + frame_count=1, + output_layout=VideoTensorLayout.bcthw, + ) + + def close(self) -> None: + self._calls.append("session.close") + + +class _Application(IApplication): + def __init__(self, calls: list[str], *, fail_to_init: bool = False) -> None: + self._calls = calls + self._fail_to_init = fail_to_init + + def init(self, commandline_args: Sequence[str]) -> None: + self._calls.append(f"application.init({list(commandline_args)!r})") + if self._fail_to_init: + raise RuntimeError("application init failed") + + def create_session(self, session_desc: SessionDesc) -> ISession: + self._calls.append("application.create_session") + return _Session(session_desc, self._calls) + + def close(self) -> None: + self._calls.append("application.close") + + +class _Window(IClientWindow): + def __init__(self, calls: list[str]) -> None: + self._calls = calls + self.results: list[StepResult] = [] + self._reported_close = False + + def get_user_input_events(self) -> UserInputEvents: + if not self._reported_close: + self._reported_close = True + return UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(0), + event_data=CloseUserInputEventData(), + ) + ] + ) + return UserInputEvents([]) + + def open(self, session_desc: SessionDesc) -> None: + del session_desc + self._calls.append("window.open") + + def write(self, result: StepResult) -> None: + self.results.append(result) + self._calls.append(f"window.write({result.step_index})") + + def close(self) -> None: + self._calls.append("window.close") + + +def _session_desc() -> SessionDesc: + return SessionDesc( + output_layout=VideoTensorLayout.bcthw, + frames_per_second_for_ui=100, + frames_per_second_for_step=30, + video_width=2, + video_height=2, + ) + + +def test_application_runner_drives_complete_lifecycle() -> None: + calls: list[str] = [] + application = _Application(calls) + window = _Window(calls) + + ApplicationRunner(application, window).run(_session_desc(), ["--model-option"]) + + assert window.results == [] + assert calls[0:3] == [ + "application.init(['--model-option'])", + "application.create_session", + "session.init", + ] + assert calls[-3:] == ["window.close", "session.close", "application.close"] + + +def test_application_runner_closes_application_when_init_fails() -> None: + calls: list[str] = [] + application = _Application(calls, fail_to_init=True) + + with pytest.raises(RuntimeError, match="application init failed"): + ApplicationRunner(application, _Window(calls)).run(_session_desc()) + + assert calls == ["application.init([])", "application.close"] diff --git a/flashdreams/test_v2/test_client_window_factory.py b/flashdreams/test_v2/test_client_window_factory.py new file mode 100644 index 000000000..ee05abe6a --- /dev/null +++ b/flashdreams/test_v2/test_client_window_factory.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the v2 client-window factory.""" + +import argparse + +import pytest + +pytestmark = pytest.mark.ci_cpu + +pytest.importorskip("aiohttp") +pytest.importorskip("aiortc") + +from flashdreams.runtime_v2.client_window_factory import create_client_window +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow + + +def test_create_client_window_selects_webrtc() -> None: + window = create_client_window( + argparse.Namespace(mode="webrtc", host="127.0.0.1", port=0) + ) + try: + assert isinstance(window, WebRTCClientWindow) + finally: + window.close() + + +def test_create_client_window_rejects_unsupported_mode() -> None: + with pytest.raises(ValueError, match="Unsupported"): + create_client_window(argparse.Namespace(mode="local")) diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py new file mode 100644 index 000000000..b7f119651 --- /dev/null +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the v2 WebRTC client window.""" + +import asyncio +import json + +import pytest +import torch + +pytestmark = pytest.mark.ci_cpu + +pytest.importorskip("aiohttp") +pytest.importorskip("aiortc") + +from aiohttp import ClientSession +from aiortc import ( + MediaStreamTrack, + RTCDataChannel, + RTCPeerConnection, + RTCSessionDescription, +) +from av import VideoFrame + +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow + + +def _session_desc() -> SessionDesc: + return SessionDesc( + output_layout=VideoTensorLayout.tchw, + frames_per_second_for_ui=30, + frames_per_second_for_step=30, + video_width=16, + video_height=16, + ) + + +async def _connect_browser( + window: WebRTCClientWindow, +) -> tuple[RTCPeerConnection, RTCDataChannel, asyncio.Future[MediaStreamTrack]]: + peer = RTCPeerConnection() + channel = peer.createDataChannel("controls") + peer.addTransceiver("video", direction="recvonly") + channel_opened = asyncio.Event() + video_track: asyncio.Future[MediaStreamTrack] = ( + asyncio.get_running_loop().create_future() + ) + + @channel.on("open") + def on_open() -> None: + channel_opened.set() + + @peer.on("track") + def on_track(track: MediaStreamTrack) -> None: + if not video_track.done(): + video_track.set_result(track) + + await peer.setLocalDescription(await peer.createOffer()) + async with ClientSession() as client: + async with client.post( + f"{window.server.url}api/webrtc/offer", + json={ + "sdp": peer.localDescription.sdp, + "type": peer.localDescription.type, + }, + ) as response: + assert response.status == 200 + answer = await response.json() + await peer.setRemoteDescription( + RTCSessionDescription(sdp=answer["sdp"], type=answer["type"]) + ) + await asyncio.wait_for(channel_opened.wait(), timeout=5) + return peer, channel, video_track + + +@pytest.mark.asyncio +async def test_window_buffers_browser_events_until_drained() -> None: + window = WebRTCClientWindow() + peer: RTCPeerConnection | None = None + try: + async with ClientSession() as client: + async with client.get(f"{window.server.url}healthz") as response: + assert response.status == 200 + assert await response.json() == { + "open": False, + "client_connected": False, + } + async with client.get(window.server.url) as response: + browser_page = await response.text() + assert response.status == 200 + assert 'id="activate"' in browser_page + assert '' in browser_page + async with client.get(f"{window.server.url}app.js") as response: + browser_script = await response.text() + assert response.status == 200 + assert 'key: "r", pressed: activationPressed' in browser_script + + window.open(_session_desc()) + peer, channel, _ = await _connect_browser(window) + channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": True})) + channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": False})) + + events = [] + for _ in range(100): + events.extend(window.get_user_input_events().get_events()) + if len(events) == 2: + break + await asyncio.sleep(0.01) + + assert len(events) == 2 + keyboard_events = [ + data + for event in events + if isinstance(data := event.get_event_data(), KeyboardUserInputEventData) + ] + assert [(event.key, event.pressed) for event in keyboard_events] == [ + ("w", True), + ("w", False), + ] + assert events[0].get_timestamp() <= events[1].get_timestamp() + assert window.get_user_input_events().get_events() == [] + finally: + if peer is not None: + await peer.close() + window.close() + + +@pytest.mark.asyncio +async def test_write_delivers_a_video_frame_to_the_browser() -> None: + window = WebRTCClientWindow() + peer: RTCPeerConnection | None = None + try: + window.open(_session_desc()) + peer, _, video_track = await _connect_browser(window) + track = await asyncio.wait_for(video_track, timeout=5) + + window.write( + StepResult( + step_index=0, + output=torch.full((2, 3, 16, 16), 17, dtype=torch.uint8), + frame_count=2, + output_layout=VideoTensorLayout.tchw, + metrics={}, + ) + ) + + frame = await asyncio.wait_for(track.recv(), timeout=5) + assert isinstance(frame, VideoFrame) + pixels = frame.to_ndarray(format="rgb24") + assert pixels.shape == (16, 16, 3) + assert abs(float(pixels.mean()) - 17.0) <= 2.0 + finally: + if peer is not None: + await peer.close() + window.close() diff --git a/integrations_v2/red_screen/README.md b/integrations_v2/red_screen/README.md index 19b1cfc04..1f31179b9 100644 --- a/integrations_v2/red_screen/README.md +++ b/integrations_v2/red_screen/README.md @@ -6,9 +6,8 @@ SPDX-License-Identifier: Apache-2.0 # Red Screen Smallest end-to-end application on the v2 API. It holds no model: a session emits -a solid red frame while the activation key is held and a solid black frame -otherwise. It runs the whole path — `IApplication`, `ISession`, `run_session`, -`IClientWindow` — on CPU. +red frames controlled by activation and intensity keys. It runs the whole path — +`IApplication`, `ISession`, `run_session`, `IClientWindow` — on CPU. ## What it demonstrates @@ -30,6 +29,28 @@ otherwise. It runs the whole path — `IApplication`, `ISession`, `run_session`, ## Usage +Start the WebRTC server and open the printed URL: + +```bash +uv run red-screen-webrtc +``` + +Hold `r` in the browser to turn the generated video red, or click **Activate** +to toggle the same keyboard event. Press `w` to increase the red intensity by +0.1 and `s` to decrease it by 0.1. Runtime options configure the browser session: + +```bash +uv run red-screen-webrtc --mode webrtc --host 0.0.0.0 --port 8080 --width 1280 --height 720 --fps 30 +``` + +Arguments after `--` belong to the application: + +```bash +uv run red-screen-webrtc -- --key x +``` + +The same application can be driven directly without WebRTC: + ```python from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.session_runner import run_session @@ -76,5 +97,6 @@ it was not asked about, which the framework tests import. Frame geometry comes from the `SessionDesc`, and the step count from the caller that drives the session. Neither is a command-line argument. -Output is a `[1, 3, 1, H, W]` float32 tensor in `bcthw` layout, with channel 0 -at `1.0` when red. +Output is a `[1, 3, 1, H, W]` float32 tensor in `bcthw` layout and the `[-1, 1]` +range expected by WebRTC. Black is `-1.0` in every channel; full red uses `1.0` +in channel 0 and `-1.0` in the other channels. diff --git a/integrations_v2/red_screen/pyproject.toml b/integrations_v2/red_screen/pyproject.toml index f55073d1f..37a2bd13e 100644 --- a/integrations_v2/red_screen/pyproject.toml +++ b/integrations_v2/red_screen/pyproject.toml @@ -13,9 +13,12 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ # Integration packages depend on the public framework, never the reverse. - "flashdreams", + "flashdreams[serving]", ] +[project.scripts] +red-screen-webrtc = "red_screen.app:main" + [tool.uv.sources] # Resolve the framework from this repository while developing the workspace. # A published integration would resolve its released FlashDreams dependency. diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 4721bc8fa..c65ac9a58 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -4,6 +4,7 @@ """Key-driven red screen application for end-to-end v2 API testing.""" import argparse +from collections.abc import Sequence from dataclasses import dataclass import torch @@ -11,11 +12,14 @@ from flashdreams.api_v2.application import IApplication from flashdreams.api_v2.session import ISession +from flashdreams.runtime_v2.application_runner import ApplicationRunner +from flashdreams.runtime_v2.client_window_factory import create_client_window from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData from flashdreams.runtime_v2.user_input_events import UserInputEvents from flashdreams.runtime_v2.video_tensor import VideoTensorLayout +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow _DEFAULT_ACTIVATION_KEY = "r" """Key that turns the screen red while held.""" @@ -36,7 +40,7 @@ class RedScreenConfig: class RedScreenSession(ISession): - """Emit a red frame while the activation key is held, black otherwise.""" + """Emit red frames controlled by activation and intensity keys.""" def __init__(self, config: RedScreenConfig, session_desc: SessionDesc) -> None: """ @@ -56,10 +60,12 @@ def __init__(self, config: RedScreenConfig, session_desc: SessionDesc) -> None: self._config = config self._session_desc = session_desc self._key_held = False + self._color_intensity = 0.0 def init(self) -> None: - """Release any held key so the session starts on a black frame.""" + """Reset key state and color intensity to start on a black frame.""" self._key_held = False + self._color_intensity = 0.0 @property def session_desc(self) -> SessionDesc: @@ -76,6 +82,10 @@ def step(self, step_index: int, events: UserInputEvents) -> StepResult: Result carrying a single ``[1, 3, 1, H, W]`` frame. """ self._apply_events(events) + import time + + # Simulate the real model inference time + time.sleep(0.1) return StepResult( step_index=step_index, output=self._frame(), @@ -88,23 +98,28 @@ def reset(self) -> None: self.init() def _apply_events(self, events: UserInputEvents) -> None: - # Events are edges, not levels: a key stays held across steps that carry - # no events for it, so only the last edge per step changes the state. - for event in events.get_events(): - data = event.get_event_data() - if ( - isinstance(data, KeyboardUserInputEventData) - and data.key == self._config.activation_key - ): - self._key_held = data.pressed + received_events = events.get_events() + if not received_events: + return + data = received_events[-1].get_event_data() + if not isinstance(data, KeyboardUserInputEventData): + return + if data.key == self._config.activation_key: + self._key_held = data.pressed + elif data.pressed and data.key.lower() == "w": + self._color_intensity = min(1.0, self._color_intensity + 0.1) + elif data.pressed and data.key.lower() == "s": + self._color_intensity = max(0.0, self._color_intensity - 0.1) def _frame(self) -> Tensor: - frame = torch.zeros( + frame = torch.full( (1, 3, 1, self._session_desc.video_height, self._session_desc.video_width), + -1.0, dtype=torch.float32, ) - if self._key_held: - frame[:, _RED_CHANNEL] = 1.0 + frame[:, _RED_CHANNEL] = ( + 1.0 if self._key_held else 2.0 * self._color_intensity - 1.0 + ) return frame @@ -112,7 +127,7 @@ def _frame(self) -> Tensor: class RedScreenApplication(IApplication): - """Application producing solid red or black frames from key input.""" + """Application producing red frames whose intensity responds to key input.""" def __init__(self) -> None: self._config: RedScreenConfig | None = None @@ -152,3 +167,61 @@ def create_session(self, session_desc: SessionDesc) -> ISession: def create_app() -> IApplication: """Return a new red screen application.""" return RedScreenApplication() + + +def _parse_args(commandline_args: Sequence[str] | None) -> argparse.Namespace: + """Parse runtime arguments and preserve application arguments.""" + parser = argparse.ArgumentParser( + prog="red-screen-webrtc", + description="Serve the red-screen v2 application.", + ) + parser.add_argument("--mode", choices=("webrtc",), default="webrtc") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=0) + parser.add_argument("--width", type=int, default=640) + parser.add_argument("--height", type=int, default=360) + parser.add_argument("--fps", type=int, default=30) + parser.add_argument( + "application_args", + nargs=argparse.REMAINDER, + help="Arguments after -- are passed to the red-screen application.", + ) + return parser.parse_args(commandline_args) + + +def main(commandline_args: Sequence[str] | None = None) -> int: + """Run red screen until the client disconnects or the process is interrupted.""" + args = _parse_args(commandline_args) + application_args = list(args.application_args) + if application_args[:1] == ["--"]: + application_args = application_args[1:] + + window = create_client_window(args) + app = create_app() + if isinstance(window, WebRTCClientWindow): + print(f"Open {window.server.url} in a browser.", flush=True) + try: + # ApplicationRunner is a FlashDreams runtime component that takes an IApplication instance, a IClientWindow instance, + # and drives the main loop. + + # TODO: in production, commandline argument parsing and IClientWindow creation should be done by flashdreams-run, a CLI tool + # basically, we need to generailze this main function to be shared by all applications + ApplicationRunner(app, window).run( + SessionDesc( + output_layout=VideoTensorLayout.bcthw, + frames_per_second_for_ui=args.fps, + frames_per_second_for_step=args.fps, + video_width=args.width, + video_height=args.height, + ), + application_args, + ) + except KeyboardInterrupt: + return 130 + finally: + window.close() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/integrations_v2/red_screen/red_screen/tests/test_red_screen.py b/integrations_v2/red_screen/red_screen/tests/test_red_screen.py index 6da51dea5..6169d160b 100644 --- a/integrations_v2/red_screen/red_screen/tests/test_red_screen.py +++ b/integrations_v2/red_screen/red_screen/tests/test_red_screen.py @@ -100,12 +100,13 @@ def _key_event(*, pressed: bool, key: str = _ACTIVATION_KEY) -> UserInputEvents: def _is_red(result: StepResult) -> bool: return bool( - torch.all(result.output[:, 0] == 1.0) and torch.all(result.output[:, 1:] == 0.0) + torch.all(result.output[:, 0] == 1.0) + and torch.all(result.output[:, 1:] == -1.0) ) def _is_black(result: StepResult) -> bool: - return bool(torch.all(result.output == 0.0)) + return bool(torch.all(result.output == -1.0)) def _run( @@ -154,6 +155,33 @@ def test_red_screen_ignores_other_keys() -> None: assert _is_black(session.step(0, _key_event(pressed=True, key="q"))) +def test_red_screen_uses_last_event_to_adjust_color_intensity() -> None: + session = _new_session() + + increased = session.step(0, _key_event(pressed=True, key="w")) + last_event_decreases = session.step( + 1, + UserInputEvents( + [ + UserInputEvent( + timestamp=uint64(0), + event_data=KeyboardUserInputEventData(key="w", pressed=True), + ), + UserInputEvent( + timestamp=uint64(1), + event_data=KeyboardUserInputEventData(key="s", pressed=True), + ), + ] + ), + ) + + assert torch.allclose( + increased.output[:, 0], + torch.full_like(increased.output[:, 0], -0.8), + ) + assert _is_black(last_event_decreases) + + def test_red_screen_starts_black_without_input() -> None: window = _run(steps=2) diff --git a/uv.lock b/uv.lock index caabe0b96..e6c769247 100644 --- a/uv.lock +++ b/uv.lock @@ -1384,11 +1384,11 @@ name = "flashdreams-red-screen" version = "0.1.0" source = { editable = "integrations_v2/red_screen" } dependencies = [ - { name = "flashdreams" }, + { name = "flashdreams", extra = ["serving"] }, ] [package.metadata] -requires-dist = [{ name = "flashdreams", editable = "flashdreams" }] +requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] [[package]] name = "flashdreams-sana-wm"