diff --git a/docs/source/developer_guides/index.rst b/docs/source/developer_guides/index.rst
index 8a8874d8b..fe2943ef0 100644
--- a/docs/source/developer_guides/index.rst
+++ b/docs/source/developer_guides/index.rst
@@ -41,6 +41,13 @@ Developer Guides
How public runner names are registered, parsed, matched to manifests,
and dispatched to integration-owned demo launch modes.
+ .. grid-item-card:: V2 WebRTC lifecycle
+ :link: v2_webrtc_lifecycle
+ :link-type: doc
+
+ Ownership and data flow across the CLI, application runner, persistent
+ browser window, resident model, and sequential sessions.
+
.. grid-item-card:: Add a new method
:link: new_integration
:link-type: doc
@@ -70,6 +77,7 @@ generated clip, see :doc:`/quickstart/index`.
inference_pipeline_overview
config_system
runner_slugs
+ v2_webrtc_lifecycle
new_integration
local_benchmarks
diff --git a/docs/source/developer_guides/v2_webrtc_lifecycle.md b/docs/source/developer_guides/v2_webrtc_lifecycle.md
new file mode 100644
index 000000000..973c75f59
--- /dev/null
+++ b/docs/source/developer_guides/v2_webrtc_lifecycle.md
@@ -0,0 +1,156 @@
+
+
+# V2 WebRTC application and session lifecycle
+
+The v2 WebRTC path keeps an application and its model alive while it creates
+one session at a time from browser requests. A session owns one rollout; it
+does not own the server or reload the model.
+
+The Self-Forcing entry command remains:
+
+```bash
+uv run --project integrations_v2/t2v_self_forcing flashdreams-run-v2 \
+ t2v-self-forcing --mode webrtc --port 8080
+```
+
+Open the printed URL, enter a prompt, and select **New session**. Closing or
+refreshing the page ends the active rollout, but the process and loaded model
+remain available for the next request. Ctrl-C ends the application.
+
+## Ownership
+
+| Component | Owns | Does not own |
+| --- | --- | --- |
+| `cli.py` | Argument parsing, application and window construction, final runner cleanup | Model state, session state, browser events |
+| `ApplicationRunner` | The initialized application and the serial application-level session loop | Browser event types or per-step scheduling |
+| `run_session()` | Exactly one session, its generation loop, its I/O thread, and session cleanup | The application or a persistent window after a successful handoff |
+| `WebRTCClientWindow` | The WebRTC server and the thread-safe browser-event queue | Application or session state |
+| `WebRTCServer` | HTTP, one active peer connection, its data channel, and its media track | Runtime lifecycle decisions |
+| `T2VApplication` | The loaded model pipeline shared by every session | A rollout cache or browser connection |
+| `T2VSession` | One prompt and one rollout cache | Model loading or server lifetime |
+
+The CLI constructs the window, then transfers its lifetime cleanup
+responsibility to the runner. While the server is idle, the runner's calling
+thread opens and polls the window. During a session, `run_session()` lends the
+window to one dedicated I/O thread, which performs the meaningful close when
+the window should end. The runner does not touch it again until that thread has
+stopped and the old session has closed. Its final, idempotent close is a fallback
+for setup failures and interrupted handoffs. This is sequential access, not
+concurrent access.
+
+## Data flow
+
+1. `cli.py` selects the `t2v-self-forcing` application and the `webrtc` window
+ mode. Arguments after `--` belong to the application.
+2. `ApplicationRunner.init()` calls `T2VApplication.init()`. The pipeline is
+ set up, moved to its device, evaluated, and retained by the application.
+3. The WebRTC mode creates `WebRTCClientWindow`, which starts its server thread
+ and exposes the browser URL.
+4. `ApplicationRunner.run()` resolves the CLI's partial `SessionDescRequest`
+ against the initialized application's default `SessionDesc`.
+5. A window whose `keeps_open_between_sessions` capability is true opens with
+ that resolved stream format before a session exists. This lets the browser
+ negotiate WebRTC and send its first request without loading a session cache
+ first. Other windows start the resolved session immediately.
+6. The server validates every data-channel message and invokes the callback
+ registered by `WebRTCClientWindow`. The callback only appends the event to a
+ thread-safe queue.
+7. While idle, `wait_for_new_session()` drains that queue. It translates the
+ latest close/new-session transition into a complete `SessionDesc` and
+ returns only that description. `ApplicationRunner` therefore never parses a
+ `UserInputEvent`.
+8. The application validates the description and creates a `T2VSession` over
+ its resident pipeline. `T2VSession.init()` creates only the per-rollout
+ prompt/cache state.
+9. `run_session()` starts one I/O thread. That thread opens the same window for
+ the actual session boundary, drains input each UI tick, calls `step_ui()`,
+ and writes completed results. The calling thread runs model steps.
+10. A new-session event stops that rollout and returns the requested next
+ `SessionDesc`. A close or natural completion returns no replacement. In
+ a persistent window the runner keeps it open, closes the old session, and
+ either starts the replacement or waits for another browser request.
+11. Ctrl-C closes the peer/server and then the application. Releasing the
+ application drops the one resident pipeline after every session cache has
+ already been released.
+
+## Why the WebRTC window is opened twice
+
+The two calls mark different boundaries:
+
+- The first call prepares the fixed stream format so a browser can connect
+ while no session exists.
+- The per-session call discards source frames queued by the previous rollout.
+ Event timestamps retain one window-lifetime monotonic origin, so opening a
+ replacement cannot reorder buffered events from the old and new browser.
+
+They do not create two servers or two peer connections. A connected peer is
+reused. The application must either accept the already-resolved `SessionDesc`
+or reject it; it cannot silently change width, height, layout, or playback rate
+after the browser stream was prepared.
+
+## Lifecycle event ordering
+
+Close and new-session events can arrive in one drained batch during a page
+refresh. Their order is meaningful, so the latest transition wins:
+
+- close, then new session: start the new browser's request;
+- new session, then close: cancel the request because that browser went away.
+
+Any close or replacement also advances the session generation. A model step
+that was already running may finish, but its result is tagged with the old
+generation and is not written. The WebRTC media track similarly clears frames
+that it has not yet handed to the encoder. At most one already-encoded frame
+can still be in flight; a transport cannot recall a frame it has sent.
+
+## WebRTC connection lifecycle
+
+Offer negotiation is serialized, and only one browser is admitted. Data
+channel and peer callbacks capture the peer that registered them, so callbacks
+from an old page cannot close or mutate its replacement. Closing the active
+data channel releases the peer immediately, allowing a refreshed page to
+negotiate without an indefinite series of HTTP 409 responses.
+
+The new-session button starts enabled and reads **Opening...** while signaling
+is in progress. A click during that interval keeps the latest valid prompt in
+the page and sends it as soon as the data channel opens. The UI therefore does
+not make model loading or a brief reconnect look like an unavailable action.
+
+The media track uses `frames_per_second_for_step`, the generated video's
+playback rate. `frames_per_second_for_ui` controls only how often the runtime
+polls input and presents completed work.
+
+The media source queue is currently unbounded. T2V rollouts are finite, and a
+session replacement clears frames that have not reached the encoder, so its
+size remains bounded by a rollout in the current application. A future
+long-running producer will need an explicit block-or-drop policy chosen for
+that application's latency requirements.
+
+## Why these boundaries are useful
+
+- The core generation pipeline and checkpoint load once and stay resident
+ across browser refreshes and sequential prompts. The existing Wan pipeline
+ may release and reload its one-shot text encoder after prompt encoding to fit
+ within GPU memory; it does not reload the diffusion model between sessions.
+- Session cleanup is complete before the next session is created, so rollout
+ caches cannot overlap accidentally.
+- WebRTC code transports validated events but never makes application
+ lifecycle decisions.
+- `ApplicationRunner.run()` visibly owns the possibly multi-session run, while
+ `run_session()` visibly owns exactly one rollout.
+- Failure paths keep the first useful exception while still releasing the
+ session I/O thread, WebRTC server thread, window, and partially initialized
+ model.
+
+`SessionDescRequest` remains separate from `SessionDesc` on purpose: the former
+means “only the fields the CLI explicitly supplied,” while the latter is the
+fully resolved contract shared by the application, session, and window. No
+additional host, session-runner class, or WebRTC-specific lifecycle wrapper is
+needed.
+
+The bundled browser page is currently prompt-oriented because text-to-video is
+the application that needs client-created sessions today. A generic UI schema
+should be introduced only when another application has a concrete, different
+request format.
diff --git a/flashdreams/flashdreams/api_v2/application.py b/flashdreams/flashdreams/api_v2/application.py
index 6a3c88002..4bde6b0f8 100644
--- a/flashdreams/flashdreams/api_v2/application.py
+++ b/flashdreams/flashdreams/api_v2/application.py
@@ -27,17 +27,19 @@ def init(self, commandline_args: Sequence[str]) -> None:
"""Parse application arguments and validate startup state."""
...
- def session_desc(self) -> SessionDesc | None:
- """Return the description of a session this application would generate.
+ def default_session_desc(self) -> SessionDesc | None:
+ """Return this initialized application's default session description.
- A caller has to describe a session before there is one to describe, and
- only the application knows what its model was trained for. Asked before
- :meth:`init`, so describing a session costs nothing.
+ The application owns its model's output requirements and defaults, such
+ as its layout, preferred dimensions, and frame rate. The runtime asks
+ after :meth:`init`, then applies the caller's explicit requests to this
+ default before calling :meth:`create_session`. That method accepts the
+ resolved description or rejects it; it does not silently change the
+ stream the runtime already asked the client window to prepare for.
Returns:
- The session to create when nobody asks for another, or ``None``,
- the default, from an application that generates whatever it is
- asked for. Its caller describes the session instead.
+ The default session description, or ``None`` when the application
+ has no output requirements and accepts the runtime's defaults.
"""
return None
@@ -49,8 +51,7 @@ def create_session(self, session_desc: SessionDesc) -> ISession:
session_desc: Session the runtime is asking for.
Returns:
- A session for ``session_desc``, resolved to what this application can
- actually produce.
+ A session that produces ``session_desc``.
Raises:
ValueError: The application cannot honour ``session_desc``.
diff --git a/flashdreams/flashdreams/api_v2/client_window.py b/flashdreams/flashdreams/api_v2/client_window.py
index 8befe1c02..b617d6d22 100644
--- a/flashdreams/flashdreams/api_v2/client_window.py
+++ b/flashdreams/flashdreams/api_v2/client_window.py
@@ -3,7 +3,7 @@
"""Client window abstract interface."""
-from abc import ABC
+from abc import ABC, abstractmethod
from .input_source import InputSource
from .output_sink import OutputSink
@@ -13,14 +13,37 @@ class IClientWindow(InputSource, OutputSink, ABC):
"""Handle application input and output for one client window.
The runtime opens the window with the session's description, then reads input
- and writes results until the run ends, and closes it then. A window stays open
- across a session reset.
+ and writes results until the run ends. A window stays open across a session
+ reset. When the client asks for a replacement session, the runtime closes the
+ old session and opens the same window with the replacement's description. A
+ persistent window can also remain open between sessions while the runtime
+ waits for another client request.
A window does not describe the output shape. The session does, and the window
is given that description in :meth:`OutputSink.open`.
- One thread makes every call on a window, so an implementation needs no
- locking except when its backend delivers input from another thread.
+ One runtime thread at a time makes every call on a window, so an
+ implementation needs no locking except when its backend delivers input from
+ another thread.
Created by the runtime, never by an application.
"""
+
+ keeps_open_between_sessions: bool = False
+ """Whether sessions start on demand and the window persists between them.
+
+ When false, the runtime starts the resolved initial session immediately and
+ returns after it ends unless the client requested a replacement. When true,
+ the runtime opens the window before creating a session, waits for a client
+ request, and returns to waiting after completion or disconnection.
+ """
+
+ @abstractmethod
+ def close(self) -> None:
+ """Release this window's resources.
+
+ This must be safe before :meth:`OutputSink.open` and after an earlier
+ call. The session loop performs the meaningful close on its I/O thread;
+ the application runner calls it again as a lifetime-cleanup fallback.
+ """
+ ...
diff --git a/flashdreams/flashdreams/api_v2/session.py b/flashdreams/flashdreams/api_v2/session.py
index 9b91a1e74..c328363c9 100644
--- a/flashdreams/flashdreams/api_v2/session.py
+++ b/flashdreams/flashdreams/api_v2/session.py
@@ -26,9 +26,11 @@ class ISession(ABC):
@abstractmethod
def init(self) -> None:
- """Load the model and anything else this run needs.
+ """Prepare the state owned by this run.
- Must not do client I/O, since this can run before a client connects.
+ Shared model loading belongs to the application. This method prepares
+ per-session state such as an encoded prompt or KV cache. It must not do
+ client I/O, since this can run before a client connects.
"""
...
diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py
index 883362abe..3b108ff7b 100644
--- a/flashdreams/flashdreams/runtime_v2/application_runner.py
+++ b/flashdreams/flashdreams/runtime_v2/application_runner.py
@@ -9,68 +9,124 @@
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
+from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest
+from flashdreams.runtime_v2.session_runner import run_session, wait_for_new_session
_LOGGER = logging.getLogger(__name__)
"""Logger for an application or window that could not be closed."""
class ApplicationRunner:
- """Create and run one application session against one client window."""
+ """Hold one initialized application and run its sessions."""
- def __init__(self, application: IApplication, client_window: IClientWindow) -> None:
+ def __init__(self, application: IApplication) -> None:
"""
Args:
- application: Long-lived application that creates the session.
- client_window: Window that supplies input and presents generated output.
+ application: Long-lived application that creates sessions.
"""
self._application = application
- self._client_window = client_window
+ self._initialized = False
+ self._initialization_attempted = False
+ self._closed = False
+
+ def init(self, commandline_args: Sequence[str] = ()) -> None:
+ """Initialize the application once.
+
+ Args:
+ commandline_args: Arguments owned and parsed by the application.
+
+ Raises:
+ RuntimeError: The application has already been initialized or closed.
+ """
+ if self._closed:
+ raise RuntimeError("ApplicationRunner is closed.")
+ if self._initialization_attempted:
+ raise RuntimeError("ApplicationRunner is already initialized.")
+ self._initialization_attempted = True
+ self._application.init(commandline_args)
+ self._initialized = True
def run(
- self, session_desc: SessionDesc, commandline_args: Sequence[str] = ()
+ self,
+ session_desc_request: SessionDescRequest,
+ client_window: IClientWindow,
) -> None:
- """Initialize the application, create one session, and run it.
-
- The run ends when the window reports a close or the session reports that
- it has finished.
+ """Create sessions against ``client_window`` until the run ends.
- The application is closed before this method returns or raises.
+ The window decides whether to start immediately with the resolved
+ description or stay open and wait for a client request. A replacement
+ description returned by the session loop starts another session after
+ the current one has closed. A persistent window returns to waiting when
+ a session finishes or its client disconnects.
- The window is closed too when the run never starts, since ``run_session``
- is what otherwise owns it, and a window may already be serving a client
- before the application has loaded anything.
+ The session and window are closed before this method returns or raises.
+ The application remains initialized, so callers can run another session
+ without reloading its shared state. Call :meth:`close` when no further
+ sessions are needed.
Args:
- session_desc: Output shape and timing requested for the session.
- commandline_args: Arguments owned and parsed by the application.
+ session_desc_request: Explicit overrides to apply to the
+ application's initialized default description.
+ client_window: Window that supplies input and presents generated output.
"""
- run_started = False
try:
- self._application.init(commandline_args)
- session = self._application.create_session(session_desc)
- run_started = True
- run_session(session, self._client_window)
+ persistent_window = client_window.keeps_open_between_sessions
+ current_session_desc = self._resolve_session_desc(session_desc_request)
+ if persistent_window:
+ client_window.open(current_session_desc)
+ next_session_desc = wait_for_new_session(
+ client_window, current_session_desc
+ )
+ else:
+ next_session_desc = current_session_desc
+
+ while next_session_desc is not None:
+ session = self._application.create_session(next_session_desc)
+ current_session_desc = session.session_desc
+ next_session_desc = run_session(
+ session,
+ client_window,
+ keep_window_open=persistent_window,
+ )
+ if persistent_window and next_session_desc is None:
+ next_session_desc = wait_for_new_session(
+ client_window, current_session_desc
+ )
finally:
- if not run_started:
- _close_client_window(self._client_window)
- _close_application(
- self._application, run_failed=sys.exc_info()[0] is not None
- )
+ _close_client_window(client_window)
+
+ def _resolve_session_desc(
+ self, session_desc_request: SessionDescRequest
+ ) -> SessionDesc:
+ """Resolve one request against the initialized application's default."""
+ if self._closed:
+ raise RuntimeError("ApplicationRunner is closed.")
+ if not self._initialized:
+ raise RuntimeError("ApplicationRunner.init() must run first.")
+ default = self._application.default_session_desc() or SessionDesc()
+ return session_desc_request.resolve(default)
+
+ def close(self) -> None:
+ """Release the application and the state it shares across sessions."""
+ if self._closed:
+ return
+ self._closed = True
+ _close_application(self._application, run_failed=sys.exc_info()[0] is not None)
def _close_client_window(client_window: IClientWindow) -> None:
- """Close a window the run never reached, so what it was serving goes with it.
+ """Ensure a window closes without hiding the run's meaningful result.
- The run has already failed by the time this is called, so a failure here is
- logged rather than raised over the top of it.
+ ``IClientWindow.close`` is idempotent. The session loop normally performs
+ the first close on its I/O thread so file-finalization errors can fail the
+ run. This is the application runner's fallback for setup failures,
+ persistent windows, and interrupted ownership handoffs.
"""
try:
client_window.close()
except Exception:
_LOGGER.exception(
- "The client window failed to close after a run that never started."
+ "The client window failed to close while the runner was stopping."
)
diff --git a/flashdreams/flashdreams/runtime_v2/cli.py b/flashdreams/flashdreams/runtime_v2/cli.py
index 438c019cb..b958aa042 100644
--- a/flashdreams/flashdreams/runtime_v2/cli.py
+++ b/flashdreams/flashdreams/runtime_v2/cli.py
@@ -5,8 +5,8 @@
``flashdreams-run-v2`` finds an application by slug, gives it the arguments
after ``--``, and hands it to :class:`ApplicationRunner` along with the window
-``--mode`` asked for. The session it asks for is the one the application says it
-would generate, with whatever the frame arguments here override.
+``--mode`` asked for. The runner initializes the application, then resolves its
+session default with whatever frame arguments this command overrides.
What the modes are, and what each one takes, belongs to
:mod:`flashdreams.runtime_v2.client_window_factory`. Nothing here reads an
@@ -14,12 +14,10 @@
"""
import argparse
+import logging
import sys
from collections.abc import Sequence
-from dataclasses import replace
-from typing import Any
-from flashdreams.api_v2.application import IApplication
from flashdreams.runtime_v2.application_registry import (
create_application,
registered_application_slugs,
@@ -29,7 +27,7 @@
add_client_window_arguments,
client_window_mode,
)
-from flashdreams.runtime_v2.session_desc import SessionDesc
+from flashdreams.runtime_v2.session_desc import SessionDescRequest
from flashdreams.runtime_v2.video_tensor import VideoTensorLayout
_ARGUMENT_SEPARATOR = "--"
@@ -39,6 +37,8 @@
command also has, so the split is stated rather than guessed.
"""
+_LOGGER = logging.getLogger(__name__)
+
def entrypoint(argv: Sequence[str] | None = None) -> None:
"""Run the command, reporting where to watch what it generates."""
@@ -55,13 +55,28 @@ def entrypoint(argv: Sequence[str] | None = None) -> None:
# Before the window, so a slug this cannot run costs nothing to find out.
application = create_application(parsed.slug)
- session_desc = _session_desc(application, parsed)
- window = mode.create(parsed)
- _report(mode.starting(window))
- # Nothing here says how long the run is: a session reports itself finished,
- # and a window ends the run when its client goes away.
- ApplicationRunner(application, window).run(session_desc, application_args)
- _report(mode.finished(window))
+ runner = ApplicationRunner(application)
+ try:
+ runner.init(application_args)
+ window = mode.create(parsed)
+ try:
+ _report(mode.starting(window))
+ except BaseException:
+ try:
+ window.close()
+ except Exception:
+ _LOGGER.exception(
+ "The client window failed to close after startup failed."
+ )
+ raise
+ # Nothing here says how long a session or window lives: each reports that.
+ try:
+ runner.run(_session_desc_request(parsed), window)
+ except KeyboardInterrupt:
+ return
+ _report(mode.finished(window))
+ finally:
+ runner.close()
def split_arguments(arguments: Sequence[str]) -> tuple[list[str], list[str]]:
@@ -132,25 +147,11 @@ def _add_session_arguments(parser: argparse.ArgumentParser) -> None:
)
-def _session_desc(
- application: IApplication, parsed_args: argparse.Namespace
-) -> SessionDesc:
- """Return the session to ask for: the application's, with the arguments on top.
-
- An application describing no session of its own gets the arguments alone,
- over :class:`SessionDesc`'s own defaults.
- """
- asked_for: dict[str, Any] = {
- field: value
- for field, value in (
- ("output_layout", parsed_args.layout),
- ("frames_per_second_for_step", parsed_args.fps),
- ("video_width", parsed_args.pixel_width),
- ("video_height", parsed_args.pixel_height),
- )
- if value is not None
- }
- described = application.session_desc()
- if described is None:
- return SessionDesc(**asked_for)
- return replace(described, **asked_for)
+def _session_desc_request(parsed_args: argparse.Namespace) -> SessionDescRequest:
+ """Return the output properties the command line explicitly requested."""
+ return SessionDescRequest(
+ output_layout=parsed_args.layout,
+ frames_per_second_for_step=parsed_args.fps,
+ video_width=parsed_args.pixel_width,
+ video_height=parsed_args.pixel_height,
+ )
diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py
index 155354296..9e5edf400 100644
--- a/flashdreams/flashdreams/runtime_v2/client_window_factory.py
+++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py
@@ -109,12 +109,16 @@ def create(self, parsed_args: argparse.Namespace) -> IClientWindow:
# Imported here so a run writing a file needs none of the serving stack.
from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow
- return WebRTCClientWindow(host=parsed_args.host, port=parsed_args.port)
+ return WebRTCClientWindow(
+ host=parsed_args.host,
+ port=parsed_args.port,
+ keeps_open_between_sessions=True,
+ )
def starting(self, client_window: IClientWindow) -> str | None:
"""Return where to connect, which nobody can guess when the port is free."""
- server = cast("WebRTCClientWindow", client_window).server
- return f"Open {server.url} in a browser."
+ window = cast("WebRTCClientWindow", client_window)
+ return f"Open {window.url} in a browser."
_MODES: tuple[ClientWindowMode, ...] = (_Mp4Mode(), _WebRTCMode())
diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/app.js b/flashdreams/flashdreams/runtime_v2/serving/web/app.js
index 029e714e7..8ed145b0f 100644
--- a/flashdreams/flashdreams/runtime_v2/serving/web/app.js
+++ b/flashdreams/flashdreams/runtime_v2/serving/web/app.js
@@ -4,11 +4,9 @@
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 newSessionButton = document.getElementById("new-session");
+const promptInput = document.getElementById("prompt");
+let pendingNewSession = null;
const send = payload => {
if (controls.readyState === "open") {
@@ -16,11 +14,36 @@ const send = payload => {
}
};
+controls.addEventListener("open", () => {
+ newSessionButton.disabled = false;
+ if (pendingNewSession !== null) {
+ send(pendingNewSession);
+ pendingNewSession = null;
+ }
+ newSessionButton.textContent = "New session";
+});
+
+controls.addEventListener("close", () => {
+ newSessionButton.disabled = true;
+ newSessionButton.textContent = "Disconnected";
+});
+
+peer.ontrack = event => {
+ document.getElementById("video").srcObject =
+ event.streams[0] ?? new MediaStream([event.track]);
+};
+
window.addEventListener("keydown", event => {
+ if (event.target === promptInput) {
+ return;
+ }
send({type: "keyboard", key: event.key, pressed: true});
});
window.addEventListener("keyup", event => {
+ if (event.target === promptInput) {
+ return;
+ }
send({type: "keyboard", key: event.key, pressed: false});
});
@@ -36,6 +59,22 @@ document.getElementById("reset").onclick = () => {
send({type: "reset"});
};
+newSessionButton.onclick = () => {
+ if (!promptInput.reportValidity()) {
+ return;
+ }
+ const request = {
+ type: "new_session",
+ metadata: {prompt: promptInput.value},
+ };
+ if (controls.readyState === "open") {
+ send(request);
+ } else {
+ pendingNewSession = request;
+ newSessionButton.textContent = "Opening...";
+ }
+};
+
window.addEventListener("beforeunload", () => send({type: "close"}));
async function connect() {
@@ -47,15 +86,26 @@ async function connect() {
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),
- });
+ let response;
+ while (true) {
+ response = await fetch("/api/webrtc/offer", {
+ method: "POST",
+ headers: {"content-type": "application/json"},
+ body: JSON.stringify(peer.localDescription),
+ });
+ if (response.status !== 409) {
+ break;
+ }
+ await new Promise(resolve => setTimeout(resolve, 100));
+ }
if (!response.ok) {
throw new Error(await response.text());
}
await peer.setRemoteDescription(await response.json());
}
-connect();
+connect().catch(error => {
+ newSessionButton.disabled = true;
+ newSessionButton.textContent = "Connection failed";
+ console.error(error);
+});
diff --git a/flashdreams/flashdreams/runtime_v2/serving/web/index.html b/flashdreams/flashdreams/runtime_v2/serving/web/index.html
index 2fb07c96e..bd37b26c1 100644
--- a/flashdreams/flashdreams/runtime_v2/serving/web/index.html
+++ b/flashdreams/flashdreams/runtime_v2/serving/web/index.html
@@ -11,6 +11,9 @@
+
+
+
diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py
index 68041879d..4e89190ec 100644
--- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py
+++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py
@@ -7,6 +7,7 @@
import asyncio
import json
+import logging
import socket
import threading
import time
@@ -18,7 +19,12 @@
import numpy as np
import torch
from aiohttp import web
-from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription
+from aiortc import (
+ MediaStreamTrack,
+ RTCDataChannel,
+ RTCPeerConnection,
+ RTCSessionDescription,
+)
from aiortc.mediastreams import MediaStreamError
from av import VideoFrame
@@ -27,6 +33,7 @@
from flashdreams.runtime_v2.user_input_event import (
CloseUserInputEventData,
KeyboardUserInputEventData,
+ NewSessionUserInputEventData,
ResetUserInputEventData,
UserInputEvent,
)
@@ -36,6 +43,8 @@
_BROWSER_PAGE = _WEB_RESOURCES.joinpath("index.html").read_text(encoding="utf-8")
_BROWSER_SCRIPT = _WEB_RESOURCES.joinpath("app.js").read_text(encoding="utf-8")
+_LOGGER = logging.getLogger(__name__)
+
class _VideoTrack(MediaStreamTrack):
"""Video track whose frames are supplied by the server."""
@@ -46,9 +55,10 @@ 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._frames: asyncio.Queue[
+ tuple[int, np.ndarray[Any, np.dtype[np.uint8]]] | None
+ ] = asyncio.Queue()
+ self._session_generation = 0
self._next_frame_time: float | None = None
self._pts = 0
self._closed = False
@@ -60,29 +70,46 @@ async def enqueue(
if self._closed:
return
for frame in frames:
- await self._frames.put(frame)
+ await self._frames.put((self._session_generation, frame))
+
+ async def start_session(self) -> None:
+ """Discard frames queued by the session being replaced."""
+ self._session_generation += 1
+ self._next_frame_time = None
+ while True:
+ try:
+ self._frames.get_nowait()
+ except asyncio.QueueEmpty:
+ return
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
+ while True:
+ item = await self._frames.get()
+ if item is None:
+ raise MediaStreamError
+ session_generation, frame = item
+ if session_generation != self._session_generation:
+ continue
+
+ 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))
+ if session_generation != self._session_generation:
+ self._next_frame_time = None
+ continue
+
+ 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."""
@@ -94,7 +121,11 @@ async def close(self) -> None:
class WebRTCServer:
- """Own the HTTP, signaling, input buffering, and media transport."""
+ """Own HTTP signaling and media transport for one browser at a time.
+
+ Validated browser events leave through the registered callback. The client
+ window owns the thread-safe queue that buffers them for the runtime loop.
+ """
def __init__(
self,
@@ -128,10 +159,12 @@ def __init__(
self._startup_error: BaseException | None = None
self._loop: asyncio.AbstractEventLoop | None = None
self._runner: web.AppRunner | None = None
+ self._offer_lock: asyncio.Lock | None = None
self._peer_connection: RTCPeerConnection | None = None
+ self._control_channel: RTCDataChannel | None = None
self._video_track: _VideoTrack | None = None
self._session_desc: SessionDesc | None = None
- self._session_start_ns: int | None = None
+ self._event_origin_ns: int | None = None
self._closed = False
self._client_connected = False
self._thread = threading.Thread(
@@ -141,6 +174,8 @@ def __init__(
)
self._thread.start()
if not self._started.wait(startup_timeout_seconds):
+ self._closed = True
+ self._stop_server_thread()
raise TimeoutError("WebRTC server did not start before the timeout.")
if self._startup_error is not None:
raise RuntimeError(
@@ -163,22 +198,42 @@ def url(self) -> str:
return f"http://{self._host}:{self._port}/"
def open(self, session_desc: SessionDesc) -> None:
- """Configure the server for one session's generated video.
+ """Configure the video format this server's sessions generate.
Args:
session_desc: Resolved dimensions, frame rate, and tensor layout.
+ A session-serving runner calls this before its first session so the
+ browser can connect and request one. Each session opens it again while
+ the browser remains connected. Its stream format must match the first
+ call because the existing media track keeps its negotiated frame rate.
+
Raises:
- RuntimeError: The server is closed or already open.
+ RuntimeError: The server is closed or has no input callback.
+ ValueError: A replacement changes the active stream format.
"""
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.")
+ current_desc = self._session_desc
+ if current_desc is not None and not _same_stream_format(
+ current_desc, session_desc
+ ):
+ raise ValueError(
+ "A replacement WebRTC session must keep the original output "
+ "layout, frame rate, width, and height."
+ )
self._session_desc = session_desc
- self._session_start_ns = time.monotonic_ns()
+ if self._event_origin_ns is None:
+ # UserInputEvents sorts by timestamp. Keep one origin for this
+ # window's lifetime so opening a replacement session cannot reorder
+ # an older browser's close after a newer browser's request.
+ self._event_origin_ns = time.monotonic_ns()
+ track = self._video_track
+ loop = self._loop
+ if track is not None and loop is not None:
+ asyncio.run_coroutine_threadsafe(track.start_session(), loop).result()
def register_input_callback(
self, callback: Callable[[UserInputEvent], None]
@@ -227,8 +282,24 @@ def close(self) -> None:
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)
+ try:
+ future.result(timeout=self._startup_timeout_seconds)
+ except BaseException:
+ try:
+ self._stop_server_thread()
+ except Exception:
+ _LOGGER.exception(
+ "The WebRTC server thread also failed to stop during cleanup."
+ )
+ raise
+ else:
+ self._stop_server_thread()
+
+ def _stop_server_thread(self) -> None:
+ """Stop and join the server thread after startup or shutdown."""
+ loop = self._loop
+ if loop is not None and not loop.is_closed():
+ 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.")
@@ -253,6 +324,7 @@ def _run_server(self) -> None:
async def _start_server(self) -> None:
"""Create and bind the standalone aiohttp application."""
+ self._offer_lock = asyncio.Lock()
app = web.Application()
app.router.add_get("/", self._serve_browser)
app.router.add_get("/app.js", self._serve_browser_script)
@@ -285,7 +357,7 @@ async def _serve_browser_script(self, _: web.Request) -> web.Response:
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."""
+ """Report whether the server can negotiate video and has a client."""
return web.json_response(
{
"open": self._session_desc is not None,
@@ -300,8 +372,6 @@ async def _offer(self, request: web.Request) -> web.Response:
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()
@@ -316,54 +386,84 @@ async def _offer(self, request: web.Request) -> web.Response:
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."
+ offer_lock = self._offer_lock
+ if offer_lock is None:
+ raise web.HTTPServiceUnavailable(reason="WebRTC server is not ready.")
+ async with offer_lock:
+ existing_peer = self._peer_connection
+ if existing_peer is not None:
+ control_channel = self._control_channel
+ if existing_peer.connectionState in {
+ "failed",
+ "disconnected",
+ "closed",
+ } or (
+ control_channel is not None
+ and control_channel.readyState == "closed"
+ ):
+ await self._release_peer_connection(existing_peer)
+ else:
+ raise web.HTTPConflict(
+ reason="A WebRTC client is already connected."
+ )
+
+ peer_connection = RTCPeerConnection()
+ video_track = _VideoTrack(session_desc.frames_per_second_for_step)
+ peer_connection.addTrack(video_track)
+ self._peer_connection = peer_connection
+ self._video_track = video_track
+
+ @peer_connection.on("datachannel")
+ def on_datachannel(channel: RTCDataChannel) -> None:
+ if self._peer_connection is not peer_connection:
+ channel.close()
+ return
+ self._control_channel = channel
+ self._client_connected = True
+
+ @channel.on("message")
+ def on_message(message: Any) -> None:
+ if self._peer_connection is not peer_connection:
+ return
+ try:
+ self._buffer_browser_message(message)
+ except ValueError as error:
+ channel.send(
+ json.dumps({"type": "error", "message": str(error)})
+ )
+
+ @channel.on("close")
+ async def on_close() -> None:
+ await self._release_peer_connection(peer_connection)
+
+ @peer_connection.on("connectionstatechange")
+ async def on_connectionstatechange() -> None:
+ if peer_connection.connectionState in {
+ "failed",
+ "disconnected",
+ "closed",
+ }:
+ await self._release_peer_connection(peer_connection)
+
+ try:
+ await peer_connection.setRemoteDescription(
+ RTCSessionDescription(sdp=sdp, type=offer_type)
+ )
+ await peer_connection.setLocalDescription(
+ await peer_connection.createAnswer()
+ )
+ local_description = peer_connection.localDescription
+ if local_description is None:
+ raise web.HTTPInternalServerError(
+ reason="WebRTC peer did not create an answer."
+ )
+ except BaseException:
+ await self._release_peer_connection(peer_connection)
+ raise
+
+ return web.json_response(
+ {"sdp": local_description.sdp, "type": local_description.type}
)
- 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."""
@@ -387,11 +487,17 @@ def _buffer_browser_message(self, raw_message: object) -> None:
event_data = KeyboardUserInputEventData(key=key, pressed=pressed)
elif event_type == "reset":
event_data = ResetUserInputEventData()
+ elif event_type == "new_session":
+ metadata = payload.get("metadata")
+ if not isinstance(metadata, dict):
+ raise ValueError("New-session event requires metadata object.")
+ event_data = NewSessionUserInputEventData(metadata=metadata)
elif event_type == "close":
event_data = CloseUserInputEventData()
else:
raise ValueError(
- "Browser event type must be 'keyboard', 'reset', or 'close'."
+ "Browser event type must be 'keyboard', 'reset', 'new_session', "
+ "or 'close'."
)
self._append_event(event_data)
@@ -400,25 +506,26 @@ def _append_event(
event_data: (
KeyboardUserInputEventData
| ResetUserInputEventData
+ | NewSessionUserInputEventData
| CloseUserInputEventData
),
) -> None:
"""Timestamp and buffer one validated browser event."""
- session_start_ns = self._session_start_ns
- if session_start_ns is None:
+ event_origin_ns = self._event_origin_ns
+ if event_origin_ns is None:
return
- timestamp_us = np.uint64((time.monotonic_ns() - session_start_ns) // 1_000)
+ timestamp_us = np.uint64((time.monotonic_ns() - event_origin_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.
+ # Every event takes the same callback path into WebRTCClientWindow's
+ # thread-safe queue; the server loop never handles lifecycle events.
callback(event)
- def _record_client_disconnect(self) -> None:
+ def _record_client_disconnect(self, peer_connection: RTCPeerConnection) -> None:
"""Buffer one close event when the active browser disconnects."""
- if not self._client_connected:
+ if self._peer_connection is not peer_connection or not self._client_connected:
return
self._client_connected = False
if not self._closed:
@@ -432,22 +539,48 @@ async def _enqueue_frames(
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
+ async def _release_peer_connection(
+ self, peer_connection: RTCPeerConnection
+ ) -> None:
+ """Release one disconnected peer so a refreshed page can reconnect."""
+ if self._peer_connection is not peer_connection:
+ return
+ self._record_client_disconnect(peer_connection)
self._peer_connection = None
+ self._control_channel = None
track = self._video_track
self._video_track = None
if track is not None:
await track.close()
- if peer_connection is not None:
+ if peer_connection.connectionState != "closed":
await peer_connection.close()
+
+ async def _shutdown(self) -> None:
+ """Release async server resources on their owning loop."""
+ peer_connection = self._peer_connection
+ if peer_connection is not None:
+ await self._release_peer_connection(peer_connection)
+ else:
+ track = self._video_track
+ self._video_track = None
+ if track is not None:
+ await track.close()
runner = self._runner
self._runner = None
if runner is not None:
await runner.cleanup()
+def _same_stream_format(first: SessionDesc, second: SessionDesc) -> bool:
+ """Return whether two sessions can use the same WebRTC media track."""
+ return (
+ first.output_layout is second.output_layout
+ and first.frames_per_second_for_step == second.frames_per_second_for_step
+ and first.video_width == second.video_width
+ and first.video_height == second.video_height
+ )
+
+
def _result_to_rgb_frames(
result: StepResult, session_desc: SessionDesc
) -> tuple[np.ndarray[Any, np.dtype[np.uint8]], ...]:
diff --git a/flashdreams/flashdreams/runtime_v2/session_desc.py b/flashdreams/flashdreams/runtime_v2/session_desc.py
index efeb7dbed..076359e4c 100644
--- a/flashdreams/flashdreams/runtime_v2/session_desc.py
+++ b/flashdreams/flashdreams/runtime_v2/session_desc.py
@@ -4,7 +4,7 @@
"""Description of the session a runtime asks an application for."""
import math
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
from typing import Any
from flashdreams.runtime_v2.video_tensor import VideoTensorLayout
@@ -14,9 +14,9 @@
class SessionDesc:
"""Description of a session, passed to create one and to open a window on it.
- The runtime fills this in to ask an application for a session, and the
- session reports back what it resolved to. The same description then
- configures the client window through ``OutputSink.open``.
+ The runtime resolves this before asking an application for a session. The
+ application either accepts it or raises; the same description configures
+ the client window through ``OutputSink.open``.
"""
output_layout: VideoTensorLayout = VideoTensorLayout.tchw
@@ -26,7 +26,7 @@ class SessionDesc:
"""Rate to read input and present finished results at, in frames per second."""
frames_per_second_for_step: int = 30
- """Rate to generate at, in frames per second. Nothing paces by it yet."""
+ """Playback rate of generated frames, in frames per second."""
video_width: int = 1280
"""Output video width in pixels."""
@@ -56,3 +56,48 @@ def __post_init__(self) -> None:
raise ValueError("SessionDesc.video_width must be > 0 when set.")
if self.video_height <= 0:
raise ValueError("SessionDesc.video_height must be > 0 when set.")
+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class SessionDescRequest:
+ """The parts of a session description a caller explicitly requests.
+
+ Leaving a field unset preserves the application's initialized default. The
+ runtime resolves this request before it asks the application to create a
+ session.
+ """
+
+ output_layout: VideoTensorLayout | None = None
+ frames_per_second_for_ui: int | None = None
+ frames_per_second_for_step: int | None = None
+ video_width: int | None = None
+ video_height: int | None = None
+ metadata: dict[str, Any] | None = None
+
+ def resolve(self, default: SessionDesc) -> SessionDesc:
+ """Return ``default`` with every explicitly requested field replaced."""
+ return replace(
+ default,
+ output_layout=(
+ default.output_layout
+ if self.output_layout is None
+ else self.output_layout
+ ),
+ frames_per_second_for_ui=(
+ default.frames_per_second_for_ui
+ if self.frames_per_second_for_ui is None
+ else self.frames_per_second_for_ui
+ ),
+ frames_per_second_for_step=(
+ default.frames_per_second_for_step
+ if self.frames_per_second_for_step is None
+ else self.frames_per_second_for_step
+ ),
+ video_width=(
+ default.video_width if self.video_width is None else self.video_width
+ ),
+ video_height=(
+ default.video_height if self.video_height is None else self.video_height
+ ),
+ metadata=default.metadata if self.metadata is None else self.metadata,
+ )
diff --git a/flashdreams/flashdreams/runtime_v2/session_runner.py b/flashdreams/flashdreams/runtime_v2/session_runner.py
index 5d013930b..016645016 100644
--- a/flashdreams/flashdreams/runtime_v2/session_runner.py
+++ b/flashdreams/flashdreams/runtime_v2/session_runner.py
@@ -7,14 +7,18 @@
import queue
import sys
import threading
+import time
+from dataclasses import replace
from enum import Enum
from flashdreams.api_v2.client_window import IClientWindow
from flashdreams.api_v2.session import ISession
from flashdreams.api_v2.user_input_event_data import UserInputEventData
+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,
+ NewSessionUserInputEventData,
ResetUserInputEventData,
UserInputEvent,
)
@@ -41,6 +45,60 @@ def _contains(events: UserInputEvents, event_type: type[UserInputEventData]) ->
)
+def _latest_session_transition(
+ events: UserInputEvents, current: SessionDesc
+) -> tuple[bool, SessionDesc | None]:
+ """Return the latest close or replacement requested by ``events``.
+
+ A new-session request replaces application-owned metadata while preserving
+ the resolved video and timing settings of the current session. The
+ application interprets and validates the metadata when it creates the next
+ session. A close after that request cancels it; a request after a close wins.
+
+ Returns:
+ Whether the batch contains a close or replacement, and the replacement
+ description. The description is ``None`` when close is latest.
+ """
+ found = False
+ next_session_desc = None
+ for event in events.get_events():
+ event_data = event.get_event_data()
+ if isinstance(event_data, CloseUserInputEventData):
+ found = True
+ next_session_desc = None
+ elif isinstance(event_data, NewSessionUserInputEventData):
+ found = True
+ next_session_desc = replace(current, metadata=dict(event_data.metadata))
+ return found, next_session_desc
+
+
+def wait_for_new_session(
+ window: IClientWindow, current_session_desc: SessionDesc
+) -> SessionDesc:
+ """Wait until ``window`` requests a new session.
+
+ This polls the window at its configured UI rate and converts the latest
+ buffered new-session event into a complete session description. Other
+ events, including a browser close, are ignored: no session exists to receive
+ them, and a persistent server must remain ready for a refreshed page.
+
+ Args:
+ window: Open client window supplying buffered input events.
+ current_session_desc: Resolved settings to preserve for the next session.
+
+ Returns:
+ A complete description for the requested session.
+ """
+ tick_seconds = 1.0 / current_session_desc.frames_per_second_for_ui
+ while True:
+ _, next_session_desc = _latest_session_transition(
+ window.get_user_input_events(), current_session_desc
+ )
+ if next_session_desc is not None:
+ return next_session_desc
+ time.sleep(tick_seconds)
+
+
def _close_session(session: ISession, *, run_failed: bool) -> None:
"""Close a session, keeping its close from hiding an earlier failure.
@@ -54,7 +112,7 @@ def _close_session(session: ISession, *, run_failed: bool) -> None:
"""
try:
session.close()
- except Exception:
+ except BaseException:
if not run_failed:
raise
_LOGGER.exception(
@@ -69,24 +127,28 @@ def run_session(
steps: int | None = None,
max_pending: int = 2,
when_full: WhenFull = WhenFull.BLOCK,
-) -> None:
+ keep_window_open: bool = False,
+) -> SessionDesc | None:
"""Drive one session against one client window.
Runs on two threads. The calling thread initializes the session and calls
``step`` for each index, with the input collected since the previous step. A
second thread owns the window: it opens it, ticks at
``frames_per_second_for_ui`` to read input, call ``step_ui`` and write
- whatever generation has finished, then closes it. A slow step therefore does
- not hold up input or output. Only the I/O thread touches the window, which is
- what a native window needs, and the window and session are always closed,
- including on failure.
+ whatever generation has finished, then closes it unless a replacement was
+ requested or ``keep_window_open`` is set. A slow step therefore does not hold
+ up input or output. Only the I/O thread touches the window, which is what a
+ native window needs. The session is always closed, and a failure always
+ closes the window.
The window ends the run by reporting a :class:`CloseUserInputEventData`, and
restarts it by reporting a :class:`ResetUserInputEventData`, which resets the
- session and takes the step index back to zero. The window stays open. Nothing
- from the abandoned generation is presented: each result carries the generation
- it was produced for, so results already waiting and a step that was still
- running when the reset arrived are both dropped rather than written.
+ session and takes the step index back to zero. A
+ :class:`NewSessionUserInputEventData` instead ends and cleans up this session,
+ leaves the window open, and becomes the complete session description returned
+ for its replacement. Nothing from an abandoned generation is presented: each
+ result carries the generation it was produced for, so results already waiting
+ and a step that was still running when the request arrived are both dropped.
Input is not split at a reset: the batch carrying it reaches the first step
afterwards whole, earlier events included. Events are edges, so a key held
@@ -120,6 +182,13 @@ def run_session(
max_pending: How many finished results may wait to be written.
when_full: What to do with a result when ``max_pending`` are already
waiting.
+ keep_window_open: Leave the window available after normal completion or
+ a client close, so an application runner can wait for another session.
+
+ Returns:
+ The complete description requested for the next session, or ``None``
+ when the run ended without one. The window remains open when a
+ description is returned or ``keep_window_open`` is set.
Raises:
ValueError: ``steps`` is negative, or ``max_pending`` is not positive.
@@ -138,9 +207,15 @@ def run_session(
try:
session.init()
- except Exception:
+ except BaseException:
# A partly initialized session still holds whatever it managed to load.
_close_session(session, run_failed=True)
+ try:
+ window.close()
+ except Exception:
+ _LOGGER.exception(
+ "The window failed to close after session initialization failed."
+ )
raise
# Backpressure is all here. Finished results wait here for the I/O thread to
@@ -157,44 +232,52 @@ def run_session(
collected_events_lock = threading.Lock()
opened = threading.Event()
stop = threading.Event()
+ io_stopped = threading.Event()
+ finish_io = threading.Event()
+ window_stays_open = False
io_failure: list[Exception] = []
+ next_session_desc: SessionDesc | None = None
# What never reached the window, reported once the run is over.
dropped_for_space = 0
- discarded_at_reset = 0
+ discarded_at_restart = 0
def present_pending_results() -> None:
"""Write every waiting result to the window, oldest first.
Because each tick writes all of them, results only pile up when writing
itself is slower than generation, not merely because the UI rate is lower.
- Results the client reset away from are dropped here rather than written,
- which is also what frees the room they were holding.
+ Results the client reset or replaced away from are dropped here rather
+ than written, which is also what frees the room they were holding.
"""
- nonlocal discarded_at_reset
+ nonlocal discarded_at_restart
while True:
try:
result_generation, result = pending_results.get_nowait()
except queue.Empty:
return
if result_generation != generation:
- discarded_at_reset += 1
+ discarded_at_restart += 1
continue
window.write(result)
def tick() -> None:
- nonlocal generation
+ nonlocal generation, next_session_desc
events = window.get_user_input_events()
+ transition_found, requested_session_desc = _latest_session_transition(
+ events, session.session_desc
+ )
with collected_events_lock:
collected_events.extend(events.get_events())
# Move on to the next generation from here, since this thread sees the
- # reset first. Under the lock, so a step already picking up its input
- # either belongs to the generation being abandoned or to the new one,
- # never to neither.
- if _contains(events, ResetUserInputEventData):
+ # reset or replacement first. Under the lock, so a step already
+ # picking up its input either belongs to the generation being
+ # abandoned or to the new one, never to neither.
+ if _contains(events, ResetUserInputEventData) or transition_found:
generation += 1
# Stop from here rather than waiting for the step loop to notice, so a
# slow step does not delay a client that has gone away.
- if _contains(events, CloseUserInputEventData):
+ if transition_found:
+ next_session_desc = requested_session_desc
stop.set()
session.step_ui(events)
present_pending_results()
@@ -217,14 +300,17 @@ def run_io() -> None:
io_failure.append(error)
finally:
opened.set()
- try:
- window.close()
- except Exception as error:
- # Closing is where a sink finishes the writes it was holding, so
- # swallowing this would report a run as complete when the output
- # never landed. An open that raised part way through gets closed
- # here too, since it still holds whatever it had acquired.
- io_failure.append(error)
+ io_stopped.set()
+ finish_io.wait()
+ if not window_stays_open:
+ try:
+ window.close()
+ except Exception as error:
+ # Closing is where a sink finishes the writes it was holding,
+ # so swallowing this would report a run as complete when the
+ # output never landed. An open that raised part way through
+ # gets closed here too, since it may hold acquired resources.
+ io_failure.append(error)
def take_collected_events() -> tuple[UserInputEvents, int]:
"""Take the input waiting for the next step, and the generation it is for."""
@@ -273,8 +359,20 @@ def add_pending_result(result_generation: int, result: StepResult) -> int:
continue
return 0
- io_thread = threading.Thread(target=run_io, name="flashdreams-io")
- io_thread.start()
+ try:
+ io_thread = threading.Thread(target=run_io, name="flashdreams-io")
+ io_thread.start()
+ except BaseException:
+ # Starting the I/O thread is part of starting the session. Nothing else
+ # can own either object when startup fails this early.
+ _close_session(session, run_failed=True)
+ try:
+ window.close()
+ except Exception:
+ _LOGGER.exception(
+ "The window failed to close after I/O thread startup failed."
+ )
+ raise
try:
opened.wait()
step_index = 0
@@ -301,29 +399,72 @@ def add_pending_result(result_generation: int, result: StepResult) -> int:
steps_run += 1
finally:
stop.set()
- io_thread.join()
+ run_failed = False
+ session_close_attempted = False
+ session_close_error: Exception | None = None
+ try:
+ try:
+ io_stopped.wait()
+ # Another session can own the same window, but only after this
+ # one has released everything it holds. If cleanup fails, tell
+ # the I/O thread to close the window and report the failure.
+ run_failed = sys.exc_info()[0] is not None
+ if (
+ (next_session_desc is not None or keep_window_open)
+ and not run_failed
+ and not io_failure
+ ):
+ session_close_attempted = True
+ try:
+ session.close()
+ except Exception as error:
+ session_close_error = error
+ else:
+ window_stays_open = True
+ finally:
+ # The I/O thread waits for the cleanup decision before it can
+ # close the window. Always release it, including when cleanup
+ # is interrupted.
+ finish_io.set()
+ io_thread.join()
+ finally:
+ # An interrupt while waiting for the I/O thread skips the normal
+ # cleanup below. Close the session here so that path cannot leak it.
+ if not session_close_attempted:
+ _close_session(
+ session,
+ run_failed=sys.exc_info()[0] is not None or bool(io_failure),
+ )
# A failure here is what the run reports, since a window failure stops
# generation rather than raising through it: both places this thread can
# be sitting give up once io_failure is set, so a run that reports a
# window failure got there without failing itself. The two are only ever
# both set by failing independently, and then this is the one raised.
- run_failed = sys.exc_info()[0] is not None
if io_failure and run_failed:
_LOGGER.error(
"The window failed as well as the run, and this is that failure.",
exc_info=io_failure[0],
)
- _close_session(session, run_failed=run_failed or bool(io_failure))
+ if session_close_error is not None and io_failure:
+ _LOGGER.error(
+ "The window failed as well as the session cleanup, and this is "
+ "that failure.",
+ exc_info=io_failure[0],
+ )
+ if session_close_error is not None:
+ raise session_close_error
# A log line is the only report of these: a caller cannot count them.
if dropped_for_space:
_LOGGER.warning(
"Dropped %d results the window could not keep up with.", dropped_for_space
)
- if discarded_at_reset:
+ if discarded_at_restart:
_LOGGER.info(
- "Discarded %d results generated before a reset.", discarded_at_reset
+ "Discarded %d results generated before a reset or replacement.",
+ discarded_at_restart,
)
if io_failure:
raise io_failure[0]
+ return next_session_desc
diff --git a/flashdreams/flashdreams/runtime_v2/user_input_event.py b/flashdreams/flashdreams/runtime_v2/user_input_event.py
index fdbd4edad..64fe9acb3 100644
--- a/flashdreams/flashdreams/runtime_v2/user_input_event.py
+++ b/flashdreams/flashdreams/runtime_v2/user_input_event.py
@@ -4,7 +4,7 @@
"""User input events, each a timestamp plus the data for one input modality."""
from dataclasses import dataclass
-from typing import ClassVar
+from typing import Any
from numpy import uint64
@@ -67,6 +67,24 @@ def get_type_name(cls) -> str:
return "reset"
+@dataclass(frozen=True, slots=True, eq=False)
+class NewSessionUserInputEventData(UserInputEventData):
+ """The client asked the application to replace the current session.
+
+ The runtime carries ``metadata`` to the application without interpreting
+ it. The application owns the requirements for its sessions and validates
+ the values when it creates the replacement.
+ """
+
+ metadata: dict[str, Any]
+ """Application-specific values requested for the replacement session."""
+
+ @classmethod
+ def get_type_name(cls) -> str:
+ """Return the event type name."""
+ return "new_session"
+
+
# Below are stubbed input event data implementations for the sake of future implementation.
@dataclass(frozen=True, slots=True, eq=False)
class MouseUserInputEventData(UserInputEventData):
@@ -133,7 +151,7 @@ class UserInputEvent:
"""User input event."""
timestamp: uint64
- """Timestamp in microseconds since the start of the session."""
+ """Microseconds since the input source's stable monotonic time origin."""
event_data: UserInputEventData
"""Event data."""
diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py
index 28936d1fc..df4387dcb 100644
--- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py
+++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py
@@ -16,12 +16,16 @@
class WebRTCClientWindow(IClientWindow):
"""Implement ``IClientWindow`` with WebRTC input and presentation."""
+ keeps_open_between_sessions = False
+ """Default direct construction to one immediate session."""
+
def __init__(
self,
*,
host: str = "127.0.0.1",
port: int = 0,
startup_timeout_seconds: float = 10.0,
+ keeps_open_between_sessions: bool = False,
) -> None:
"""Create the WebRTC backend.
@@ -32,9 +36,12 @@ def __init__(
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.
+ keeps_open_between_sessions: Wait for browser-requested sessions and
+ keep serving after each one. False runs one session immediately.
"""
+ self.keeps_open_between_sessions = keeps_open_between_sessions
self._input_events: queue.SimpleQueue[UserInputEvent] = queue.SimpleQueue()
- self.server = WebRTCServer(
+ self._server = WebRTCServer(
host=host,
port=port,
startup_timeout_seconds=startup_timeout_seconds,
@@ -44,15 +51,20 @@ 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)
+ self._server.register_input_callback(handle_input)
+
+ @property
+ def url(self) -> str:
+ """Return the URL at which a browser can open this window."""
+ return self._server.url
def open(self, session_desc: SessionDesc) -> None:
- """Implement ``OutputSink.open`` by configuring WebRTC output.
+ """Configure WebRTC output for waiting or running a session.
Args:
session_desc: Resolved dimensions, frame rate, and tensor layout.
"""
- self.server.open(session_desc)
+ self._server.open(session_desc)
def get_user_input_events(self) -> UserInputEvents:
"""Implement ``InputSource.get_user_input_events`` for browser input.
@@ -73,8 +85,8 @@ def write(self, result: StepResult) -> None:
Args:
result: Generated frames matching the opened session.
"""
- self.server.write(result)
+ self._server.write(result)
def close(self) -> None:
"""Implement ``OutputSink.close`` by releasing WebRTC resources."""
- self.server.close()
+ self._server.close()
diff --git a/flashdreams/flashdreams/t2v_v2/application.py b/flashdreams/flashdreams/t2v_v2/application.py
index f21ad8364..0060b59bb 100644
--- a/flashdreams/flashdreams/t2v_v2/application.py
+++ b/flashdreams/flashdreams/t2v_v2/application.py
@@ -23,8 +23,8 @@
class T2VSessionConfig:
"""What one command line resolved to, shared by every session it creates."""
- prompt: str
- """Text every session generates from."""
+ prompt: str | None
+ """Default text for a session, when its request does not provide one."""
device: str
"""Device the pipeline is built on."""
@@ -41,8 +41,9 @@ class T2VApplication(IApplication):
command line for all of them. An integration supplies
:class:`T2VApplicationDefaults` and inherits the rest.
- The model is loaded once, on the first session, and shared by every session
- after it, since loading reads a checkpoint of several gigabytes.
+ :meth:`init` loads the model once after resolving its command-line options.
+ The application then keeps that model resident and shares it with every
+ session, since loading reads a checkpoint of several gigabytes.
"""
session_type: type[T2VSession] = T2VSession
@@ -65,21 +66,27 @@ def pipeline_config(self) -> Any:
return self._pipeline_config
def init(self, commandline_args: Sequence[str]) -> None:
- """Parse what to generate, how much of it, and where.
+ """Parse application options and load the shared model.
Not what size or rate to generate at: that describes the session, which
- the caller asks for. The model is not loaded here either.
+ the caller asks for. Loading happens after device, compilation, seed,
+ and integration-specific options have been resolved.
Raises:
- ValueError: No prompt was given, or the rollout length is not one
- this model can generate.
+ ValueError: An explicitly provided prompt is empty, or the rollout
+ length is not one this model can generate.
"""
parser = argparse.ArgumentParser(
prog="flashdreams-run-v2 SLUG --",
description="Generate video from text.",
)
parser.add_argument(
- "--prompt", default="", help="Text to generate from. Required."
+ "--prompt",
+ default=None,
+ help=(
+ "Default text to generate from. A client may instead provide "
+ "a prompt for each session."
+ ),
)
parser.add_argument(
"--device",
@@ -117,8 +124,8 @@ def init(self, commandline_args: Sequence[str]) -> None:
self._configure_argument_parser(parser)
args = parser.parse_args(list(commandline_args))
- if not args.prompt.strip():
- raise ValueError("--prompt is required, and cannot be empty.")
+ if args.prompt is not None and not args.prompt.strip():
+ raise ValueError("--prompt cannot be empty.")
self._validate_total_blocks(args.total_blocks)
self._apply_parsed_arguments(args)
@@ -130,13 +137,22 @@ def init(self, commandline_args: Sequence[str]) -> None:
self._pipeline_config = self._apply_seed_override(
self._pipeline_config, args.seed
)
- self._config = T2VSessionConfig(
+ config = T2VSessionConfig(
prompt=args.prompt,
device=args.device,
total_blocks=args.total_blocks,
)
-
- def session_desc(self) -> SessionDesc:
+ # Take ownership as soon as setup returns. If moving or evaluating the
+ # model fails, ApplicationRunner.close() can still release what loaded.
+ pipeline = self._pipeline_config.setup()
+ self._pipeline = pipeline
+ pipeline = pipeline.to(config.device)
+ self._pipeline = pipeline
+ pipeline = pipeline.eval()
+ self._config = config
+ self._pipeline = pipeline
+
+ def default_session_desc(self) -> SessionDesc:
"""Return the description of a session this application uses.
A model generates its best video at the size and rate it was trained
@@ -151,26 +167,28 @@ def session_desc(self) -> SessionDesc:
)
def create_session(self, session_desc: SessionDesc) -> ISession:
- """Create one uninitialized session, loading the model if needed.
+ """Create one uninitialized session against the resident model.
Raises:
RuntimeError: :meth:`init` has not run yet.
- ValueError: The description asks for output this cannot generate.
+ ValueError: The description asks for output this cannot generate,
+ or its prompt metadata is not a non-empty string.
"""
config = self._config
- if config is None:
+ pipeline = self._pipeline
+ if config is None or pipeline is None:
raise RuntimeError(
f"{type(self).__name__}.init() must run before create_session()."
)
- # Before loading rather than after: a checkpoint of several gigabytes is
- # a long wait for a layout this was never going to accept.
self._validate_layout(session_desc)
- if self._pipeline is None:
- self._pipeline = self._pipeline_config.setup().to(config.device).eval()
- self._validate_frame_size(session_desc, self._pipeline)
- return self.session_type(
- self._pipeline, config.prompt, session_desc, config.total_blocks
- )
+ prompt = session_desc.metadata.get("prompt", config.prompt)
+ if not isinstance(prompt, str) or not prompt.strip():
+ raise ValueError(
+ "A session prompt is required: pass --prompt or set "
+ "SessionDesc.metadata['prompt'] to a non-empty string."
+ )
+ self._validate_frame_size(session_desc, pipeline)
+ return self.session_type(pipeline, prompt, session_desc, config.total_blocks)
def close(self) -> None:
"""Release the model, and whatever memory it was holding."""
diff --git a/flashdreams/flashdreams/t2v_v2/testing.py b/flashdreams/flashdreams/t2v_v2/testing.py
index 4751ee907..9c545ed56 100644
--- a/flashdreams/flashdreams/t2v_v2/testing.py
+++ b/flashdreams/flashdreams/t2v_v2/testing.py
@@ -132,7 +132,7 @@ def check_t2v_model_impl(
application.init(commandline_args)
try:
if session_desc is None:
- session_desc = application.session_desc()
+ session_desc = application.default_session_desc()
run_session(
application.create_session(session_desc),
_InspectingClientWindow(inspector),
diff --git a/flashdreams/test_v2/test_application_runner.py b/flashdreams/test_v2/test_application_runner.py
index e77b73732..783658ea9 100644
--- a/flashdreams/test_v2/test_application_runner.py
+++ b/flashdreams/test_v2/test_application_runner.py
@@ -4,6 +4,7 @@
"""CPU tests for the v2 application runner."""
import logging
+import threading
from collections.abc import Sequence
import pytest
@@ -13,11 +14,13 @@
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 import application_runner as application_runner_module
from flashdreams.runtime_v2.application_runner import ApplicationRunner
-from flashdreams.runtime_v2.session_desc import SessionDesc
+from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest
from flashdreams.runtime_v2.step_result import StepResult
from flashdreams.runtime_v2.user_input_event import (
CloseUserInputEventData,
+ NewSessionUserInputEventData,
UserInputEvent,
)
from flashdreams.runtime_v2.user_input_events import UserInputEvents
@@ -76,12 +79,15 @@ def __init__(
*,
fail_to_init: bool = False,
fail_to_close: bool = False,
+ fail_to_create_at: int | None = None,
session_length: int | None = None,
) -> None:
self._calls = calls
self._fail_to_init = fail_to_init
self._fail_to_close = fail_to_close
+ self._fail_to_create_at = fail_to_create_at
self._session_length = session_length
+ self.created_session_descs: list[SessionDesc] = []
def init(self, commandline_args: Sequence[str]) -> None:
self._calls.append(f"application.init({list(commandline_args)!r})")
@@ -90,6 +96,9 @@ def init(self, commandline_args: Sequence[str]) -> None:
def create_session(self, session_desc: SessionDesc) -> ISession:
self._calls.append("application.create_session")
+ if self._fail_to_create_at == len(self.created_session_descs):
+ raise RuntimeError("session creation failed")
+ self.created_session_descs.append(session_desc)
return _Session(session_desc, self._calls, length=self._session_length)
def close(self) -> None:
@@ -103,6 +112,7 @@ def __init__(self, calls: list[str]) -> None:
self._calls = calls
self.results: list[StepResult] = []
self._reported_close = False
+ self._closed = False
def get_user_input_events(self) -> UserInputEvents:
if not self._reported_close:
@@ -126,6 +136,9 @@ def write(self, result: StepResult) -> None:
self._calls.append(f"window.write({result.step_index})")
def close(self) -> None:
+ if self._closed:
+ return
+ self._closed = True
self._calls.append("window.close")
@@ -136,8 +149,48 @@ def get_user_input_events(self) -> UserInputEvents:
return UserInputEvents([])
-def _session_desc() -> SessionDesc:
- return SessionDesc(
+class _ScriptedWindow(_Window):
+ """Report one scripted event batch each time the runner polls."""
+
+ def __init__(self, calls: list[str], events: list[UserInputEvents]) -> None:
+ super().__init__(calls)
+ self._events = list(events)
+
+ def get_user_input_events(self) -> UserInputEvents:
+ if self._events:
+ return self._events.pop(0)
+ return UserInputEvents([])
+
+
+class _ServingWindow(_Window):
+ """Request two sessions, then interrupt the persistent runner."""
+
+ keeps_open_between_sessions = True
+
+ def __init__(self, calls: list[str]) -> None:
+ super().__init__(calls)
+ self._prompts = ["A cat surfing", "A dog snowboarding"]
+
+ def get_user_input_events(self) -> UserInputEvents:
+ if threading.current_thread() is not threading.main_thread():
+ return UserInputEvents([])
+ completed_sessions = self._calls.count("session.close")
+ if completed_sessions == len(self._prompts):
+ raise KeyboardInterrupt
+ return UserInputEvents(
+ [
+ UserInputEvent(
+ timestamp=uint64(completed_sessions),
+ event_data=NewSessionUserInputEventData(
+ metadata={"prompt": self._prompts[completed_sessions]}
+ ),
+ )
+ ]
+ )
+
+
+def _session_desc_request() -> SessionDescRequest:
+ return SessionDescRequest(
output_layout=VideoTensorLayout.bcthw,
frames_per_second_for_ui=100,
frames_per_second_for_step=30,
@@ -146,45 +199,251 @@ def _session_desc() -> SessionDesc:
)
-def test_application_runner_drives_complete_lifecycle() -> None:
+def test_application_runner_keeps_the_application_open_for_another_session() -> None:
calls: list[str] = []
- application = _Application(calls)
- window = _Window(calls)
-
- ApplicationRunner(application, window).run(_session_desc(), ["--model-option"])
-
- assert window.results == []
+ application = _Application(calls, session_length=1)
+ runner = ApplicationRunner(application)
+ first_window = _SilentWindow(calls)
+ second_window = _SilentWindow(calls)
+
+ runner.init(["--model-option"])
+ runner.run(_session_desc_request(), first_window)
+ runner.run(_session_desc_request(), second_window)
+
+ assert [result.step_index for result in first_window.results] == [0]
+ assert [result.step_index for result in second_window.results] == [0]
+ assert calls.count("application.create_session") == 2
+ assert calls.count("application.close") == 0
assert calls[0:3] == [
"application.init(['--model-option'])",
"application.create_session",
"session.init",
]
- assert calls[-3:] == ["window.close", "session.close", "application.close"]
+ runner.close()
+ assert calls[-1] == "application.close"
-def test_application_runner_closes_both_when_the_run_never_starts() -> None:
- """The window is closed by the loop, which a failure here never reaches, and
- a window may already be serving a client by then."""
+def test_application_runner_replaces_a_session_from_window_metadata() -> None:
calls: list[str] = []
- application = _Application(calls, fail_to_init=True)
+ application = _Application(calls)
+ runner = ApplicationRunner(application)
+ new_session = UserInputEvents(
+ [
+ UserInputEvent(
+ timestamp=uint64(0),
+ event_data=NewSessionUserInputEventData(
+ metadata={"prompt": "A dog snowboarding"}
+ ),
+ )
+ ]
+ )
+ close = UserInputEvents(
+ [
+ UserInputEvent(
+ timestamp=uint64(1),
+ event_data=CloseUserInputEventData(),
+ )
+ ]
+ )
+ window = _ScriptedWindow(calls, [new_session, close])
+
+ runner.init()
+ runner.run(_session_desc_request(), window)
+ runner.close()
+
+ assert len(application.created_session_descs) == 2
+ assert application.created_session_descs[0].metadata == {}
+ assert application.created_session_descs[1].metadata == {
+ "prompt": "A dog snowboarding"
+ }
+ assert calls.count("session.close") == 2
+ assert calls.count("window.open") == 2
+ assert calls.count("window.close") == 1
+ first_creation = calls.index("application.create_session")
+ second_creation = calls.index("application.create_session", first_creation + 1)
+ assert calls.index("session.close") < second_creation
+
+
+def test_application_runner_keeps_a_persistent_window_between_sessions() -> None:
+ calls: list[str] = []
+ application = _Application(calls, session_length=1)
+ runner = ApplicationRunner(application)
+ window = _ServingWindow(calls)
- with pytest.raises(RuntimeError, match="application init failed"):
- ApplicationRunner(application, _Window(calls)).run(_session_desc())
+ runner.init()
+ with pytest.raises(KeyboardInterrupt):
+ runner.run(_session_desc_request(), window)
- assert calls == ["application.init([])", "window.close", "application.close"]
+ assert [desc.metadata for desc in application.created_session_descs] == [
+ {"prompt": "A cat surfing"},
+ {"prompt": "A dog snowboarding"},
+ ]
+ assert calls.count("session.close") == 2
+ assert calls.count("window.open") == 3
+ assert calls.count("window.close") == 1
+ assert calls.count("application.close") == 0
+ runner.close()
+ assert calls[-1] == "application.close"
+
+
+def test_persistent_window_returns_to_waiting_after_a_client_disconnect() -> None:
+ class DisconnectingWindow(_ScriptedWindow):
+ keeps_open_between_sessions = True
+
+ def get_user_input_events(self) -> UserInputEvents:
+ if not self._events:
+ raise KeyboardInterrupt
+ return super().get_user_input_events()
+
+ def lifecycle_event(
+ timestamp: int,
+ event_data: CloseUserInputEventData | NewSessionUserInputEventData,
+ ) -> UserInputEvents:
+ return UserInputEvents(
+ [UserInputEvent(timestamp=uint64(timestamp), event_data=event_data)]
+ )
+
+ calls: list[str] = []
+ application = _Application(calls)
+ window = DisconnectingWindow(
+ calls,
+ [
+ lifecycle_event(
+ 0, NewSessionUserInputEventData(metadata={"prompt": "first"})
+ ),
+ lifecycle_event(1, CloseUserInputEventData()),
+ lifecycle_event(
+ 2, NewSessionUserInputEventData(metadata={"prompt": "second"})
+ ),
+ lifecycle_event(3, CloseUserInputEventData()),
+ ],
+ )
+ runner = ApplicationRunner(application)
+ runner.init()
+
+ with pytest.raises(KeyboardInterrupt):
+ runner.run(_session_desc_request(), window)
+
+ assert [desc.metadata for desc in application.created_session_descs] == [
+ {"prompt": "first"},
+ {"prompt": "second"},
+ ]
+ assert calls.count("session.close") == 2
+ assert calls.count("window.open") == 3
+ assert calls.count("window.close") == 1
+ runner.close()
-def test_application_runner_ends_a_run_a_window_cannot_end() -> None:
- """A window with no client never reports a close, so the session ends it."""
+@pytest.mark.parametrize("persistent", [False, True])
+def test_application_runner_closes_the_window_when_session_creation_fails(
+ persistent: bool,
+) -> None:
calls: list[str] = []
- window = _SilentWindow(calls)
+ application = _Application(calls, fail_to_create_at=0)
+ window = _ServingWindow(calls) if persistent else _Window(calls)
+ runner = ApplicationRunner(application)
+ runner.init()
+
+ with pytest.raises(RuntimeError, match="session creation failed"):
+ runner.run(_session_desc_request(), window)
- ApplicationRunner(_Application(calls, session_length=3), window).run(
- _session_desc()
+ assert calls.count("window.close") == 1
+ runner.close()
+
+
+def test_application_runner_closes_the_window_when_replacement_creation_fails() -> None:
+ calls: list[str] = []
+ application = _Application(calls, fail_to_create_at=1)
+ window = _ScriptedWindow(
+ calls,
+ [
+ UserInputEvents(
+ [
+ UserInputEvent(
+ timestamp=uint64(0),
+ event_data=NewSessionUserInputEventData(
+ metadata={"prompt": "replacement"}
+ ),
+ )
+ ]
+ )
+ ],
)
+ runner = ApplicationRunner(application)
+ runner.init()
+
+ with pytest.raises(RuntimeError, match="session creation failed"):
+ runner.run(_session_desc_request(), window)
+
+ assert calls.count("session.close") == 1
+ assert calls.count("window.close") == 1
+ runner.close()
+
+
+def test_persistent_window_closes_if_session_handoff_is_interrupted(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ def interrupt_after_cleanup(
+ session: ISession,
+ window: IClientWindow,
+ *,
+ keep_window_open: bool,
+ ) -> SessionDesc | None:
+ del window
+ assert keep_window_open
+ session.close()
+ raise KeyboardInterrupt
+
+ monkeypatch.setattr(
+ application_runner_module, "run_session", interrupt_after_cleanup
+ )
+ calls: list[str] = []
+ runner = ApplicationRunner(_Application(calls))
+ window = _ServingWindow(calls)
+ runner.init()
+
+ with pytest.raises(KeyboardInterrupt):
+ runner.run(_session_desc_request(), window)
+
+ assert calls.count("session.close") == 1
+ assert calls.count("window.close") == 1
+ runner.close()
+
+
+def test_application_runner_closes_the_window_when_a_session_cannot_start() -> None:
+ calls: list[str] = []
+ runner = ApplicationRunner(_Application(calls))
+ window = _Window(calls)
+
+ with pytest.raises(RuntimeError, match=r"init\(\) must run first"):
+ runner.run(_session_desc_request(), window)
+
+ assert calls == ["window.close"]
+
+
+def test_serving_closes_the_window_when_the_session_request_is_invalid() -> None:
+ calls: list[str] = []
+ runner = ApplicationRunner(_Application(calls))
+ window = _ServingWindow(calls)
+ runner.init()
+
+ with pytest.raises(ValueError, match="frames_per_second_for_ui"):
+ runner.run(SessionDescRequest(frames_per_second_for_ui=0), window)
+
+ assert calls == ["application.init([])", "window.close"]
+ runner.close()
+
+
+def test_application_runner_rejects_a_second_initialization() -> None:
+ calls: list[str] = []
+ runner = ApplicationRunner(_Application(calls))
+
+ runner.init()
+ with pytest.raises(RuntimeError, match="already initialized"):
+ runner.init()
+ runner.close()
- assert [result.step_index for result in window.results] == [0, 1, 2]
- assert calls[-3:] == ["window.close", "session.close", "application.close"]
+ assert calls == ["application.init([])", "application.close"]
def test_application_runner_reports_the_run_rather_than_the_close(
@@ -192,9 +451,13 @@ def test_application_runner_reports_the_run_rather_than_the_close(
) -> None:
calls: list[str] = []
application = _Application(calls, fail_to_init=True, fail_to_close=True)
+ runner = ApplicationRunner(application)
with caplog.at_level(logging.ERROR, logger=_RUNNER_LOGGER):
with pytest.raises(RuntimeError, match="application init failed"):
- ApplicationRunner(application, _Window(calls)).run(_session_desc())
+ try:
+ runner.init()
+ finally:
+ runner.close()
assert "application close failed" in caplog.text
diff --git a/flashdreams/test_v2/test_cli.py b/flashdreams/test_v2/test_cli.py
index ff65f7e8c..2b288438b 100644
--- a/flashdreams/test_v2/test_cli.py
+++ b/flashdreams/test_v2/test_cli.py
@@ -13,7 +13,6 @@
import shutil
from collections.abc import Sequence
from pathlib import Path
-from typing import Any
import pytest
import torch
@@ -98,6 +97,28 @@ def create_session(self, session_desc: SessionDesc) -> ISession:
return OneStepSession(session_desc)
+class InitializedDescriptionApplication(IApplication):
+ """Application whose default output width comes from its own arguments."""
+
+ def __init__(self) -> None:
+ self._width: int | None = None
+ self.asked_for: SessionDesc | None = None
+
+ def init(self, commandline_args: Sequence[str]) -> None:
+ if len(commandline_args) != 2 or commandline_args[0] != "--width":
+ raise ValueError("--width is required.")
+ self._width = int(commandline_args[1])
+
+ def default_session_desc(self) -> SessionDesc:
+ if self._width is None:
+ raise RuntimeError("init() must run before default_session_desc().")
+ return SessionDesc(video_width=self._width)
+
+ def create_session(self, session_desc: SessionDesc) -> ISession:
+ self.asked_for = session_desc
+ return OneStepSession(session_desc)
+
+
class OneStepSession(ISession):
"""A session generating one frame and reporting itself finished."""
@@ -133,6 +154,7 @@ class RecordingWindow(IClientWindow):
def __init__(self) -> None:
self.results: list[StepResult] = []
+ self.closed = False
def get_user_input_events(self) -> UserInputEvents:
return UserInputEvents([])
@@ -144,7 +166,7 @@ def write(self, result: StepResult) -> None:
self.results.append(result)
def close(self) -> None:
- return
+ self.closed = True
## Splitting the command line
@@ -378,13 +400,45 @@ def test_nothing_is_measured_unless_a_run_asks(
assert list(tmp_path.glob("*.json")) == []
-def test_an_application_that_will_not_start_reports_why(
+def test_a_non_serving_run_without_a_prompt_reports_why(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_install(monkeypatch, StubT2VApplication(_stand_in()), RecordingWindow())
- with pytest.raises(ValueError, match="--prompt is required"):
- cli.entrypoint(["stub", "--mode", "webrtc"])
+ with pytest.raises(ValueError, match="session prompt is required"):
+ cli.entrypoint(["stub", "--mode", "mp4"])
+
+
+def test_a_browser_server_starts_without_a_command_line_prompt(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class InterruptingWindow(RecordingWindow):
+ keeps_open_between_sessions = True
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.closed = False
+
+ def get_user_input_events(self) -> UserInputEvents:
+ raise KeyboardInterrupt
+
+ def close(self) -> None:
+ self.closed = True
+
+ pipeline = _stand_in()
+ window = InterruptingWindow()
+ _install(monkeypatch, StubT2VApplication(pipeline))
+ monkeypatch.setattr(
+ cli,
+ "client_window_mode",
+ lambda name: StubMode(name, window),
+ )
+
+ cli.entrypoint(["stub", "--mode", "webrtc"])
+
+ assert window.closed
+ assert pipeline.device == "cpu"
+ assert pipeline.eval_count == 1
## Describing the session to run
@@ -437,6 +491,18 @@ def test_a_model_generates_what_it_was_trained_for_unless_asked_otherwise(
assert window.session_desc.video_height == pipeline.height
+def test_application_arguments_resolve_the_default_session_before_creation(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ application = InitializedDescriptionApplication()
+ _install(monkeypatch, application, RecordingWindow())
+
+ cli.entrypoint(["stub", "--mode", "webrtc", "--", "--width", "64"])
+
+ assert application.asked_for is not None
+ assert application.asked_for.video_width == 64
+
+
## The command itself
@@ -460,6 +526,31 @@ def test_the_run_goes_to_the_window_the_mode_asked_for(
assert len(window.results) == _TOTAL_BLOCKS
+def test_a_window_is_closed_when_reporting_its_start_fails(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class FailingStartMode(StubMode):
+ def starting(self, client_window: IClientWindow) -> str | None:
+ del client_window
+ raise RuntimeError("cannot report window")
+
+ pipeline = _stand_in()
+ application = StubT2VApplication(pipeline)
+ window = RecordingWindow()
+ _install(monkeypatch, application)
+ monkeypatch.setattr(
+ cli,
+ "client_window_mode",
+ lambda name: FailingStartMode(name, window),
+ )
+
+ with pytest.raises(RuntimeError, match="cannot report window"):
+ cli.entrypoint(["stub", "--mode", "webrtc"])
+
+ assert window.closed
+ assert pipeline.closed
+
+
def test_the_command_needs_somewhere_to_write() -> None:
with pytest.raises(SystemExit):
cli.entrypoint(["stub"])
diff --git a/flashdreams/test_v2/test_client_window_factory.py b/flashdreams/test_v2/test_client_window_factory.py
index b7e675cbb..a0892f1c8 100644
--- a/flashdreams/test_v2/test_client_window_factory.py
+++ b/flashdreams/test_v2/test_client_window_factory.py
@@ -36,6 +36,7 @@ def test_a_run_goes_to_a_file_unless_it_says_otherwise(tmp_path: Path) -> None:
assert parsed.mode == "mp4"
assert isinstance(window, Mp4ClientWindow)
+ assert window.keeps_open_between_sessions is False
def test_a_file_run_with_nowhere_to_write_says_so() -> None:
@@ -73,6 +74,7 @@ def test_a_browser_run_is_told_where_to_connect(self) -> None:
window = mode.create(_parsed(["--mode", "webrtc"]))
try:
assert isinstance(window, WebRTCClientWindow)
- assert mode.starting(window) == f"Open {window.server.url} in a browser."
+ assert mode.starting(window) == f"Open {window.url} in a browser."
+ assert window.keeps_open_between_sessions is True
finally:
window.close()
diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py
index 7928aa412..58bff6c69 100644
--- a/flashdreams/test_v2/test_session_runner.py
+++ b/flashdreams/test_v2/test_session_runner.py
@@ -5,6 +5,7 @@
import logging
import threading
+from dataclasses import replace
import pytest
import torch
@@ -14,11 +15,16 @@
from flashdreams.api_v2.session import ISession
from flashdreams.api_v2.user_input_event_data import UserInputEventData
from flashdreams.runtime_v2.session_desc import SessionDesc
-from flashdreams.runtime_v2.session_runner import WhenFull, run_session
+from flashdreams.runtime_v2.session_runner import (
+ WhenFull,
+ run_session,
+ wait_for_new_session,
+)
from flashdreams.runtime_v2.step_result import StepResult
from flashdreams.runtime_v2.user_input_event import (
CloseUserInputEventData,
KeyboardUserInputEventData,
+ NewSessionUserInputEventData,
ResetUserInputEventData,
UserInputEvent,
)
@@ -326,6 +332,122 @@ def test_run_session_stops_when_the_window_reports_a_close() -> None:
assert log.calls[-2:] == ["window.close", "session.close"]
+def test_run_session_returns_the_next_session_desc_after_cleanup() -> None:
+ log = CallLog()
+ session_desc = _session_desc()
+ session = FakeSession(session_desc, log)
+ request = NewSessionUserInputEventData(metadata={"prompt": "A cat surfing"})
+ window = RecordingClientWindow(log, [_lifecycle_event(request)])
+
+ returned = run_session(session, window, steps=None)
+
+ assert returned == replace(session_desc, metadata={"prompt": "A cat surfing"})
+ assert "session.close" in log.calls
+ assert "window.close" not in log.calls
+ assert "session.step(0)" not in log.calls
+
+
+def test_the_latest_close_or_new_session_event_wins() -> None:
+ session_desc = _session_desc()
+ replacement = NewSessionUserInputEventData(metadata={"prompt": "new browser"})
+
+ close_then_new = RecordingClientWindow(
+ CallLog(),
+ [
+ UserInputEvents(
+ [
+ UserInputEvent(uint64(0), CloseUserInputEventData()),
+ UserInputEvent(uint64(1), replacement),
+ ]
+ )
+ ],
+ )
+ returned = run_session(
+ FakeSession(session_desc, CallLog()), close_then_new, steps=None
+ )
+
+ assert returned == replace(session_desc, metadata={"prompt": "new browser"})
+
+ new_then_close = RecordingClientWindow(
+ CallLog(),
+ [
+ UserInputEvents(
+ [
+ UserInputEvent(uint64(0), replacement),
+ UserInputEvent(uint64(1), CloseUserInputEventData()),
+ ]
+ )
+ ],
+ )
+ returned = run_session(
+ FakeSession(session_desc, CallLog()), new_then_close, steps=None
+ )
+
+ assert returned is None
+
+
+def test_wait_for_new_session_ignores_a_request_cancelled_by_close() -> None:
+ session_desc = _session_desc()
+ replacement = NewSessionUserInputEventData(metadata={"prompt": "new browser"})
+ window = RecordingClientWindow(
+ CallLog(),
+ [
+ UserInputEvents(
+ [
+ UserInputEvent(uint64(0), replacement),
+ UserInputEvent(uint64(1), CloseUserInputEventData()),
+ ]
+ ),
+ UserInputEvents(
+ [
+ UserInputEvent(uint64(2), CloseUserInputEventData()),
+ UserInputEvent(uint64(3), replacement),
+ ]
+ ),
+ ],
+ )
+
+ returned = wait_for_new_session(window, session_desc)
+
+ assert returned == replace(session_desc, metadata={"prompt": "new browser"})
+
+
+def test_run_session_closes_the_window_when_replacement_cleanup_fails() -> None:
+ log = CallLog()
+ session = FakeSession(_session_desc(), log, fail_to_close=True)
+ window = RecordingClientWindow(
+ log,
+ [
+ _lifecycle_event(
+ NewSessionUserInputEventData(metadata={"prompt": "A cat surfing"})
+ )
+ ],
+ )
+
+ with pytest.raises(RuntimeError, match="session close failed"):
+ run_session(session, window, steps=None)
+
+ assert log.calls[-2:] == ["session.close", "window.close"]
+
+
+def test_run_session_releases_the_io_thread_when_cleanup_is_interrupted() -> None:
+ log = CallLog()
+
+ class InterruptedCleanupSession(FakeSession):
+ def close(self) -> None:
+ super().close()
+ raise KeyboardInterrupt
+
+ session = InterruptedCleanupSession(_session_desc(), log)
+ window = RecordingClientWindow(log)
+
+ with pytest.raises(KeyboardInterrupt):
+ run_session(session, window, steps=0, keep_window_open=True)
+
+ assert log.calls[-2:] == ["session.close", "window.close"]
+ assert not any(thread.name == _IO_THREAD_NAME for thread in threading.enumerate())
+
+
def test_run_session_resets_the_session_and_the_step_index() -> None:
log = CallLog()
session = FakeSession(_session_desc(), log)
@@ -354,6 +476,18 @@ def test_run_session_stops_when_the_session_says_it_has_finished() -> None:
assert [result.step_index for result in window.results] == [0, 1]
+def test_run_session_can_leave_the_window_open_after_completion() -> None:
+ log = CallLog()
+ session = FiniteSession(_session_desc(), log, length=1)
+ window = RecordingClientWindow(log)
+
+ returned = run_session(session, window, keep_window_open=True)
+
+ assert returned is None
+ assert "session.close" in log.calls
+ assert "window.close" not in log.calls
+
+
def test_run_session_ends_at_whichever_comes_first() -> None:
"""A caller can ask for fewer steps than the session would generate."""
log = CallLog()
@@ -379,7 +513,7 @@ def test_run_session_lets_a_reset_restart_a_finished_session() -> None:
assert "session.reset" in log.calls
-def test_run_session_closes_a_session_that_failed_to_init() -> None:
+def test_run_session_closes_a_session_and_window_that_failed_to_init() -> None:
log = CallLog()
class FailingSession(FakeSession):
@@ -393,9 +527,83 @@ def init(self) -> None:
with pytest.raises(RuntimeError, match="init failed"):
run_session(session, window, steps=1)
- # A session that got halfway through starting still has to be released, and
- # the window is never opened for a session that cannot run.
- assert log.calls == ["session.init", "session.close"]
+ # A session that got halfway through starting still has to be released. The
+ # constructor-owned window may already hold resources even though open was
+ # never called, so it is closed too.
+ assert log.calls == ["session.init", "session.close", "window.close"]
+
+
+def test_run_session_closes_a_session_whose_init_is_interrupted() -> None:
+ log = CallLog()
+
+ class InterruptedSession(FakeSession):
+ def init(self) -> None:
+ super().init()
+ raise KeyboardInterrupt
+
+ session = InterruptedSession(_session_desc(), log)
+
+ with pytest.raises(KeyboardInterrupt):
+ run_session(session, RecordingClientWindow(log), steps=1)
+
+ assert log.calls == ["session.init", "session.close", "window.close"]
+
+
+def test_run_session_closes_the_session_when_the_io_thread_cannot_start(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ def fail_to_start(thread: threading.Thread) -> None:
+ del thread
+ raise RuntimeError("thread start failed")
+
+ monkeypatch.setattr(threading.Thread, "start", fail_to_start)
+ log = CallLog()
+
+ with pytest.raises(RuntimeError, match="thread start failed"):
+ run_session(
+ FakeSession(_session_desc(), log), RecordingClientWindow(log), steps=1
+ )
+
+ assert log.calls == ["session.init", "session.close", "window.close"]
+
+
+def test_run_session_closes_the_session_when_io_shutdown_is_interrupted(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ real_event = threading.Event
+ events_created = 0
+
+ class InterruptingEvent(real_event):
+ """Interrupt the caller waiting for the runner's I/O-stopped event."""
+
+ def __init__(self, *, interrupt_wait: bool = False) -> None:
+ super().__init__()
+ self._interrupt_wait = interrupt_wait
+
+ def wait(self, timeout: float | None = None) -> bool:
+ if (
+ self._interrupt_wait
+ and threading.current_thread() is threading.main_thread()
+ ):
+ raise KeyboardInterrupt
+ return super().wait(timeout)
+
+ def event_factory() -> InterruptingEvent:
+ nonlocal events_created
+ events_created += 1
+ # run_session creates opened, stop, then io_stopped in this order.
+ return InterruptingEvent(interrupt_wait=events_created == 3)
+
+ monkeypatch.setattr(threading, "Event", event_factory)
+ log = CallLog()
+
+ with pytest.raises(KeyboardInterrupt):
+ run_session(
+ FakeSession(_session_desc(), log), RecordingClientWindow(log), steps=0
+ )
+
+ assert log.calls[-2:] == ["window.close", "session.close"]
+ assert not any(thread.name == _IO_THREAD_NAME for thread in threading.enumerate())
def test_run_session_gives_the_step_after_a_reset_the_whole_batch() -> None:
@@ -475,6 +683,34 @@ def get_user_input_events(self) -> UserInputEvents:
assert [result.step_index for result in window.results] == [0]
+def test_run_session_drops_a_result_generated_during_disconnect() -> None:
+ log = CallLog()
+ disconnect_reported = threading.Event()
+
+ class SlowStep(FakeSession):
+ def step(self, step_index: int, events: UserInputEvents) -> StepResult:
+ del events
+ disconnect_reported.wait()
+ return super().step(step_index, UserInputEvents([]))
+
+ class DisconnectingWindow(RecordingClientWindow):
+ def get_user_input_events(self) -> UserInputEvents:
+ events = super().get_user_input_events()
+ if events.get_events():
+ disconnect_reported.set()
+ return events
+
+ session = SlowStep(_session_desc(), log)
+ window = DisconnectingWindow(
+ log,
+ [UserInputEvents([]), _lifecycle_event(CloseUserInputEventData())],
+ )
+
+ run_session(session, window, steps=None)
+
+ assert window.results == []
+
+
def test_run_session_presents_every_result_when_blocking() -> None:
log = CallLog()
session = FakeSession(_session_desc(), log)
diff --git a/flashdreams/test_v2/test_t2v_application.py b/flashdreams/test_v2/test_t2v_application.py
index 0efebf7da..4d58cad9c 100644
--- a/flashdreams/test_v2/test_t2v_application.py
+++ b/flashdreams/test_v2/test_t2v_application.py
@@ -8,6 +8,7 @@
"""
import argparse
+from dataclasses import replace
from types import SimpleNamespace
from typing import Any
@@ -72,8 +73,8 @@ def close(self) -> None:
class FakePipelineConfig:
"""Record how often the model was loaded, which is the expensive part."""
- def __init__(self) -> None:
- self.pipeline = FakePipeline()
+ def __init__(self, pipeline: FakePipeline | None = None) -> None:
+ self.pipeline = pipeline if pipeline is not None else FakePipeline()
self.setup_count = 0
def setup(self) -> FakePipeline:
@@ -88,6 +89,7 @@ def __init__(
self, pipeline: Any, prompt: str, session_desc: SessionDesc, total_blocks: int
) -> None:
super().__init__(pipeline, prompt, session_desc, total_blocks)
+ self.prompt = prompt
self.blocks_to_generate = total_blocks
@@ -136,7 +138,7 @@ def _rollout_length(app: T2VApplication) -> int:
An application answers no such question: the length reaches a run through
the session, which reports itself finished at the end of it.
"""
- session = app.create_session(app.session_desc())
+ session = app.create_session(app.default_session_desc())
assert isinstance(session, RecordingRollout)
return session.blocks_to_generate
@@ -228,7 +230,7 @@ def test_a_model_describes_the_clip_it_was_trained_for_without_loading() -> None
describe a session before it can ask for one."""
app = T2VApplication(defaults=_defaults())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert desc == _session_desc()
assert _pipeline_config(app).setup_count == 0
@@ -240,12 +242,21 @@ def test_the_rollout_length_can_be_overridden() -> None:
assert _rollout_length(app) == 3
-def test_a_run_needs_something_to_generate_from() -> None:
+def test_an_application_preloads_before_a_session_has_a_prompt() -> None:
+ defaults = _defaults()
+ app = T2VApplication(defaults=defaults)
+
+ app.init([])
+
+ assert defaults.pipeline_config.setup_count == 1
+ with pytest.raises(ValueError, match="session prompt is required"):
+ app.create_session(_session_desc())
+
+
+def test_an_explicit_command_line_prompt_cannot_be_empty() -> None:
app = T2VApplication(defaults=_defaults())
- with pytest.raises(ValueError, match="--prompt is required"):
- app.init([])
- with pytest.raises(ValueError, match="--prompt is required"):
+ with pytest.raises(ValueError, match="--prompt cannot be empty"):
app.init(["--prompt", " "])
@@ -264,10 +275,14 @@ def test_no_session_is_created_before_the_application_is_told_what_to_do() -> No
## Loading the model
-def test_the_model_loads_once_and_every_session_shares_it() -> None:
- app = _application()
+def test_initialization_loads_the_model_once_for_every_session() -> None:
+ app = ApplicationUnderTest(defaults=_defaults())
config = _pipeline_config(app)
+ assert config.setup_count == 0
+
+ app.init(["--prompt", _PROMPT])
+
first = app.create_session(_session_desc())
second = app.create_session(_session_desc())
@@ -277,28 +292,65 @@ def test_the_model_loads_once_and_every_session_shares_it() -> None:
assert first is not second
+def test_session_metadata_can_replace_the_command_line_prompt() -> None:
+ app = _application()
+
+ session = app.create_session(
+ replace(_session_desc(), metadata={"prompt": "A dog snowboarding"})
+ )
+
+ assert isinstance(session, RecordingRollout)
+ assert session.prompt == "A dog snowboarding"
+
+
+@pytest.mark.parametrize("prompt", ["", 7])
+def test_a_replacement_prompt_must_be_non_empty_text(prompt: object) -> None:
+ app = _application()
+
+ with pytest.raises(ValueError, match="session prompt is required"):
+ app.create_session(replace(_session_desc(), metadata={"prompt": prompt}))
+
+ assert _pipeline_config(app).setup_count == 1
+
+
def test_closing_the_application_releases_the_model() -> None:
app = _application()
config = _pipeline_config(app)
- app.create_session(_session_desc())
app.close()
assert config.pipeline.closed
+def test_failed_model_initialization_remains_owned_for_cleanup() -> None:
+ class FailingPipeline(FakePipeline):
+ def eval(self) -> "FakePipeline":
+ super().eval()
+ raise RuntimeError("eval failed")
+
+ pipeline = FailingPipeline()
+ app = T2VApplication(
+ defaults=_defaults(pipeline_config=FakePipelineConfig(pipeline))
+ )
+
+ with pytest.raises(RuntimeError, match="eval failed"):
+ app.init(["--prompt", _PROMPT])
+ app.close()
+
+ assert pipeline.closed
+
+
## What a model will not generate
-def test_a_layout_the_model_does_not_emit_is_refused_before_it_loads() -> None:
- """A checkpoint of several gigabytes is a long wait for a certain refusal."""
+def test_a_layout_the_model_does_not_emit_is_refused_before_a_session_starts() -> None:
app = _application()
config = _pipeline_config(app)
with pytest.raises(ValueError, match="only produces tchw output"):
app.create_session(_session_desc(VideoTensorLayout.bcthw))
- assert config.setup_count == 0
+ assert config.setup_count == 1
@pytest.mark.parametrize("width,height", [(130, 64), (128, 60)])
@@ -408,9 +460,11 @@ def test_a_seed_reaches_the_model_where_a_model_keeps_one() -> None:
"""Straight onto the config the pipeline is built from, so a model that
draws its own noise draws the same noise twice."""
diffusion_model = SimpleNamespace(seed=42)
- defaults = _defaults(
- pipeline_config=SimpleNamespace(diffusion_model=diffusion_model)
+ pipeline_config = SimpleNamespace(
+ diffusion_model=diffusion_model,
+ setup=lambda: FakePipeline(),
)
+ defaults = _defaults(pipeline_config=pipeline_config)
app = T2VApplication(defaults=defaults)
app.init(["--prompt", _PROMPT, "--seed", "7"])
diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py
index b7f119651..2e1e299fe 100644
--- a/flashdreams/test_v2/test_webrtc_client_window.py
+++ b/flashdreams/test_v2/test_webrtc_client_window.py
@@ -5,6 +5,8 @@
import asyncio
import json
+import threading
+from dataclasses import replace
import pytest
import torch
@@ -25,7 +27,11 @@
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_event import (
+ CloseUserInputEventData,
+ KeyboardUserInputEventData,
+ NewSessionUserInputEventData,
+)
from flashdreams.runtime_v2.video_tensor import VideoTensorLayout
from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow
@@ -33,8 +39,8 @@
def _session_desc() -> SessionDesc:
return SessionDesc(
output_layout=VideoTensorLayout.tchw,
- frames_per_second_for_ui=30,
- frames_per_second_for_step=30,
+ frames_per_second_for_ui=60,
+ frames_per_second_for_step=20,
video_width=16,
video_height=16,
)
@@ -42,7 +48,11 @@ def _session_desc() -> SessionDesc:
async def _connect_browser(
window: WebRTCClientWindow,
-) -> tuple[RTCPeerConnection, RTCDataChannel, asyncio.Future[MediaStreamTrack]]:
+) -> tuple[
+ RTCPeerConnection,
+ RTCDataChannel,
+ asyncio.Future[MediaStreamTrack],
+]:
peer = RTCPeerConnection()
channel = peer.createDataChannel("controls")
peer.addTransceiver("video", direction="recvonly")
@@ -61,16 +71,23 @@ def on_track(track: MediaStreamTrack) -> None:
video_track.set_result(track)
await peer.setLocalDescription(await peer.createOffer())
+ deadline = asyncio.get_running_loop().time() + 5
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()
+ while True:
+ async with client.post(
+ f"{window.url}api/webrtc/offer",
+ json={
+ "sdp": peer.localDescription.sdp,
+ "type": peer.localDescription.type,
+ },
+ ) as response:
+ if response.status != 409:
+ assert response.status == 200
+ answer = await response.json()
+ break
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError("The previous browser did not disconnect.")
+ await asyncio.sleep(0.01)
await peer.setRemoteDescription(
RTCSessionDescription(sdp=answer["sdp"], type=answer["type"])
)
@@ -83,36 +100,56 @@ async def test_window_buffers_browser_events_until_drained() -> None:
window = WebRTCClientWindow()
peer: RTCPeerConnection | None = None
try:
+ assert window.keeps_open_between_sessions is False
async with ClientSession() as client:
- async with client.get(f"{window.server.url}healthz") as response:
+ async with client.get(f"{window.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:
+ async with client.get(window.url) as response:
browser_page = await response.text()
assert response.status == 200
assert 'id="activate"' in browser_page
+ assert 'id="prompt"' in browser_page
+ assert 'id="new-session" type="button">' in browser_page
+ assert (
+ 'id="new-session" type="button">Opening...' in browser_page
+ )
assert '' in browser_page
- async with client.get(f"{window.server.url}app.js") as response:
+ async with client.get(f"{window.url}app.js") as response:
browser_script = await response.text()
assert response.status == 200
assert 'key: "r", pressed: activationPressed' in browser_script
+ assert 'type: "new_session"' in browser_script
+ assert "metadata: {prompt: promptInput.value}" in browser_script
+ assert "pendingNewSession = request" in browser_script
+ assert 'newSessionButton.textContent = "Opening..."' in browser_script
+ assert "event.target === promptInput" in browser_script
+ assert "response.status !== 409" 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}))
+ channel.send(
+ json.dumps(
+ {
+ "type": "new_session",
+ "metadata": {"prompt": "A dog snowboarding"},
+ }
+ )
+ )
events = []
for _ in range(100):
events.extend(window.get_user_input_events().get_events())
- if len(events) == 2:
+ if len(events) == 3:
break
await asyncio.sleep(0.01)
- assert len(events) == 2
+ assert len(events) == 3
keyboard_events = [
data
for event in events
@@ -122,11 +159,161 @@ async def test_window_buffers_browser_events_until_drained() -> None:
("w", True),
("w", False),
]
+ new_session_events = [
+ data
+ for event in events
+ if isinstance(data := event.get_event_data(), NewSessionUserInputEventData)
+ ]
+ assert [event.metadata for event in new_session_events] == [
+ {"prompt": "A dog snowboarding"}
+ ]
assert events[0].get_timestamp() <= events[1].get_timestamp()
assert window.get_user_input_events().get_events() == []
+
+ # Make the old timestamp epoch observably older than a fresh offer, then
+ # reopen between the old browser's close and the refreshed request. The
+ # drained FIFO order must survive that session boundary.
+ await asyncio.sleep(1)
+ await peer.close()
+ peer = None
+ async with ClientSession() as client:
+ deadline = asyncio.get_running_loop().time() + 5
+ while True:
+ async with client.get(f"{window.url}healthz") as response:
+ health = await response.json()
+ if not health["client_connected"]:
+ break
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError("The closed browser was not released.")
+ await asyncio.sleep(0.01)
+ window.open(replace(_session_desc(), metadata={"prompt": "next"}))
+ peer, channel, _ = await _connect_browser(window)
+ channel.send(
+ json.dumps(
+ {
+ "type": "new_session",
+ "metadata": {"prompt": "A fox in a forest"},
+ }
+ )
+ )
+ refreshed_events = []
+ for _ in range(100):
+ refreshed_events.extend(window.get_user_input_events().get_events())
+ if any(
+ isinstance(event.get_event_data(), NewSessionUserInputEventData)
+ for event in refreshed_events
+ ):
+ break
+ await asyncio.sleep(0.01)
+ assert [
+ event_data.metadata
+ for event in refreshed_events
+ if isinstance(
+ event_data := event.get_event_data(),
+ NewSessionUserInputEventData,
+ )
+ ] == [{"prompt": "A fox in a forest"}]
+ lifecycle_types = [
+ type(event.get_event_data())
+ for event in refreshed_events
+ if isinstance(
+ event.get_event_data(),
+ (CloseUserInputEventData, NewSessionUserInputEventData),
+ )
+ ]
+ assert lifecycle_types == [
+ CloseUserInputEventData,
+ NewSessionUserInputEventData,
+ ]
finally:
if peer is not None:
await peer.close()
+ await asyncio.sleep(0.05)
+ window.close()
+
+
+@pytest.mark.asyncio
+async def test_overlapping_offer_connects_after_active_browser_releases() -> None:
+ window = WebRTCClientWindow()
+ first = RTCPeerConnection()
+ second = RTCPeerConnection()
+ try:
+ window.open(_session_desc())
+ for peer in (first, second):
+ peer.createDataChannel("controls")
+ peer.addTransceiver("video", direction="recvonly")
+ await peer.setLocalDescription(await peer.createOffer())
+
+ async with ClientSession() as client:
+
+ async def offer(
+ peer: RTCPeerConnection,
+ ) -> tuple[int, dict[str, str] | None]:
+ async with client.post(
+ f"{window.url}api/webrtc/offer",
+ json={
+ "sdp": peer.localDescription.sdp,
+ "type": peer.localDescription.type,
+ },
+ ) as response:
+ answer = await response.json() if response.status == 200 else None
+ return response.status, answer
+
+ responses = await asyncio.gather(offer(first), offer(second))
+
+ assert sorted(status for status, _ in responses) == [200, 409]
+ admitted_index = next(
+ index for index, (status, _) in enumerate(responses) if status == 200
+ )
+ rejected_index = 1 - admitted_index
+ peers = (first, second)
+ admitted = peers[admitted_index]
+ rejected = peers[rejected_index]
+ answer = responses[admitted_index][1]
+ assert answer is not None
+ await admitted.setRemoteDescription(
+ RTCSessionDescription(sdp=answer["sdp"], type=answer["type"])
+ )
+ deadline = asyncio.get_running_loop().time() + 5
+ while admitted.connectionState != "connected":
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError("The admitted peer did not connect.")
+ await asyncio.sleep(0.01)
+
+ # The refreshed browser overlaps the old one and must be rejected
+ # while that old peer remains active.
+ status, _ = await offer(rejected)
+ assert status == 409
+
+ # Once the old page releases its peer, retrying that same pending
+ # offer must succeed within a bound.
+ await admitted.close()
+ deadline = asyncio.get_running_loop().time() + 5
+ replacement_answer = None
+ while replacement_answer is None:
+ status, replacement_answer = await offer(rejected)
+ if status != 409:
+ assert status == 200
+ break
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError("The refreshed browser did not connect.")
+ await asyncio.sleep(0.01)
+ assert replacement_answer is not None
+ await rejected.setRemoteDescription(
+ RTCSessionDescription(
+ sdp=replacement_answer["sdp"],
+ type=replacement_answer["type"],
+ )
+ )
+ deadline = asyncio.get_running_loop().time() + 5
+ while rejected.connectionState != "connected":
+ if asyncio.get_running_loop().time() >= deadline:
+ raise TimeoutError("The refreshed browser did not connect.")
+ await asyncio.sleep(0.01)
+ finally:
+ await first.close()
+ await second.close()
+ await asyncio.sleep(0.05)
window.close()
@@ -154,7 +341,58 @@ async def test_write_delivers_a_video_frame_to_the_browser() -> None:
pixels = frame.to_ndarray(format="rgb24")
assert pixels.shape == (16, 16, 3)
assert abs(float(pixels.mean()) - 17.0) <= 2.0
+
+ window.open(replace(_session_desc(), metadata={"prompt": "replacement"}))
+ window.write(
+ StepResult(
+ step_index=0,
+ output=torch.full((2, 3, 16, 16), 211, dtype=torch.uint8),
+ frame_count=2,
+ output_layout=VideoTensorLayout.tchw,
+ metrics={},
+ )
+ )
+ # One frame may already be inside the encoder when the replacement
+ # starts. The track discards the rest of the old source queue, so the
+ # replacement must arrive immediately after that in-flight frame.
+ replacement_pixels = None
+ for _ in range(2):
+ replacement_frame = await asyncio.wait_for(track.recv(), timeout=5)
+ assert isinstance(replacement_frame, VideoFrame)
+ pixels = replacement_frame.to_ndarray(format="rgb24")
+ if float(pixels.mean()) > 100:
+ replacement_pixels = pixels
+ break
+ assert replacement_pixels is not None
+ assert abs(float(replacement_pixels.mean()) - 211.0) <= 2.0
+ assert peer.connectionState == "connected"
finally:
if peer is not None:
await peer.close()
+ await asyncio.sleep(0.05)
+ window.close()
+
+
+def test_an_open_peer_keeps_its_media_format_across_sessions() -> None:
+ window = WebRTCClientWindow()
+ try:
+ session_desc = _session_desc()
+ window.open(session_desc)
+
+ # UI polling is a runtime concern and does not change the media track.
+ window.open(replace(session_desc, frames_per_second_for_ui=30))
+ with pytest.raises(ValueError, match="original output"):
+ window.open(replace(session_desc, frames_per_second_for_step=10))
+ finally:
window.close()
+
+
+def test_closing_a_window_stops_its_server_thread() -> None:
+ window = WebRTCClientWindow()
+
+ window.close()
+ window.close()
+
+ assert not any(
+ thread.name == "flashdreams-webrtc" for thread in threading.enumerate()
+ )
diff --git a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py
index dc7e4acc3..eaf46acf0 100644
--- a/integrations_v2/color_fade/color_fade/tests/test_color_fade.py
+++ b/integrations_v2/color_fade/color_fade/tests/test_color_fade.py
@@ -16,7 +16,7 @@
from flashdreams.api_v2.session import ISession
from flashdreams.runtime_v2.application_runner import ApplicationRunner
from flashdreams.runtime_v2.mp4_client_window import Mp4ClientWindow
-from flashdreams.runtime_v2.session_desc import SessionDesc
+from flashdreams.runtime_v2.session_desc import SessionDesc, SessionDescRequest
from flashdreams.runtime_v2.step_result import StepResult
from flashdreams.runtime_v2.user_input_events import UserInputEvents
from flashdreams.runtime_v2.video_tensor import VideoTensorLayout
@@ -260,10 +260,22 @@ def test_a_run_writes_the_whole_fade_to_an_mp4(tmp_path: Path) -> None:
frames_per_step = 5
# No step count: the session knows how long its fade is and ends the run.
- ApplicationRunner(create_app(), Mp4ClientWindow(path)).run(
- _session_desc(width=_PLAYABLE_WIDTH, height=_PLAYABLE_HEIGHT),
- ["--seconds", str(_SECONDS), "--frames-per-step", str(frames_per_step)],
- )
+ runner = ApplicationRunner(create_app())
+ try:
+ runner.init(
+ ["--seconds", str(_SECONDS), "--frames-per-step", str(frames_per_step)]
+ )
+ runner.run(
+ SessionDescRequest(
+ output_layout=VideoTensorLayout.bcthw,
+ frames_per_second_for_step=_FRAMES_PER_SECOND,
+ video_width=_PLAYABLE_WIDTH,
+ video_height=_PLAYABLE_HEIGHT,
+ ),
+ Mp4ClientWindow(path),
+ )
+ finally:
+ runner.close()
frames = _decode(path, width=_PLAYABLE_WIDTH, height=_PLAYABLE_HEIGHT)
assert len(frames) == _STEPS_FOR_THE_FADE * frames_per_step
diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py
index 7a4e64892..620f92c79 100644
--- a/integrations_v2/red_screen/red_screen/app.py
+++ b/integrations_v2/red_screen/red_screen/app.py
@@ -13,8 +13,7 @@
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.session_desc import SessionDesc, SessionDescRequest
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
@@ -203,30 +202,32 @@ def main(commandline_args: Sequence[str] | None = None) -> int:
if application_args[:1] == ["--"]:
application_args = application_args[1:]
- window = create_client_window(args)
+ window = WebRTCClientWindow(
+ host=args.host,
+ port=args.port,
+ keeps_open_between_sessions=False,
+ )
app = create_app()
+ runner = ApplicationRunner(app)
if isinstance(window, WebRTCClientWindow):
- print(f"Open {window.server.url} in a browser.", flush=True)
+ print(f"Open {window.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(
+ runner.init(application_args)
+ runner.run(
+ SessionDescRequest(
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,
+ window,
)
except KeyboardInterrupt:
return 130
finally:
window.close()
+ runner.close()
return 0
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 d6f666764..c4d30f503 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
@@ -8,6 +8,7 @@
import pytest
import torch
from numpy import uint64
+from red_screen import app as red_screen_app
from red_screen import create_app
from flashdreams.api_v2.client_window import IClientWindow
@@ -16,6 +17,7 @@
from flashdreams.runtime_v2.session_runner import WhenFull, run_session
from flashdreams.runtime_v2.step_result import StepResult
from flashdreams.runtime_v2.user_input_event import (
+ CloseUserInputEventData,
KeyboardUserInputEventData,
UserInputEvent,
)
@@ -135,6 +137,63 @@ def _new_session() -> ISession:
## Tests
+def test_webrtc_entrypoint_runs_its_initial_session(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class EntrypointWindow(IClientWindow):
+ def __init__(
+ self,
+ *,
+ host: str,
+ port: int,
+ keeps_open_between_sessions: bool,
+ ) -> None:
+ assert host == "127.0.0.1"
+ assert port == 8080
+ assert keeps_open_between_sessions is False
+ self.keeps_open_between_sessions = keeps_open_between_sessions
+ self.url = "http://127.0.0.1:8080"
+ self.opened = False
+ self.closed = False
+ self._polled = False
+ windows.append(self)
+
+ def get_user_input_events(self) -> UserInputEvents:
+ if self._polled:
+ raise RuntimeError("entrypoint waited for a session request")
+ self._polled = True
+ return UserInputEvents(
+ [
+ UserInputEvent(
+ timestamp=uint64(0),
+ event_data=CloseUserInputEventData(),
+ )
+ ]
+ )
+
+ def open(self, session_desc: SessionDesc) -> None:
+ del session_desc
+ self.opened = True
+
+ def write(self, result: StepResult) -> None:
+ del result
+
+ def close(self) -> None:
+ self.closed = True
+
+ windows: list[EntrypointWindow] = []
+ monkeypatch.setattr(red_screen_app, "WebRTCClientWindow", EntrypointWindow)
+
+ exit_code = red_screen_app.main(
+ ["--port", "8080", "--width", "2", "--height", "2", "--fps", "100"]
+ )
+
+ assert exit_code == 0
+ assert len(windows) == 1
+ assert windows[0].opened
+ assert windows[0].closed
+
+
def test_red_screen_holds_red_between_key_edges() -> None:
# Key down at step 0 and up at step 2. The step in between carries no events,
# so it exercises held state rather than a repeated key-down.
diff --git a/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py
index 2ef24d233..80d8f7f76 100644
--- a/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py
+++ b/integrations_v2/t2v_causal_forcing/t2v_causal_forcing/tests/test_stand_in_model.py
@@ -26,7 +26,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
integration already ships rather than written down again."""
app = CausalForcingT2VApplication(pipeline_config=FakeT2VPipelineConfig())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert (desc.video_width, desc.video_height) == (
RUNNER_WAN21_T2V_1PT3B_CHUNKWISE.pixel_width,
diff --git a/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py b/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py
index 303dd5e35..31fc223fa 100644
--- a/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py
+++ b/integrations_v2/t2v_cosmos_predict2/t2v_cosmos_predict2/tests/test_stand_in_model.py
@@ -26,7 +26,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
integration already ships rather than written down again."""
app = CosmosPredict2T2VApplication(pipeline_config=FakeT2VPipelineConfig())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert (desc.video_width, desc.video_height) == (
RUNNER_COSMOS2_T2V_2B_720P.pixel_width,
diff --git a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py
index 669aba346..5c7697210 100644
--- a/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py
+++ b/integrations_v2/t2v_fastvideo_causal_wan22/t2v_fastvideo_causal_wan22/tests/test_stand_in_model.py
@@ -8,12 +8,15 @@
``test_real_model.py``.
"""
+import copy
+from typing import Any
+
import pytest
from fastvideo_causal_wan22.config import RUNNER_WAN22_T2V_14B
from t2v_fastvideo_causal_wan22 import FastvideoCausalWan22T2VApplication
from flashdreams.runtime_v2.video_tensor import VideoTensorLayout
-from flashdreams.t2v_v2.testing import FakeT2VPipelineConfig
+from flashdreams.t2v_v2.testing import FakeT2VPipeline, FakeT2VPipelineConfig
pytestmark = pytest.mark.ci_cpu
@@ -26,7 +29,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
integration already ships rather than written down again."""
app = FastvideoCausalWan22T2VApplication(pipeline_config=FakeT2VPipelineConfig())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert (desc.video_width, desc.video_height) == (
RUNNER_WAN22_T2V_14B.pixel_width,
@@ -44,7 +47,13 @@ def test_compilation_is_turned_off_for_both_noise_level_transformers() -> None:
model splits denoising across two transformers, and the shared override
reaches only one of them, so it is overridden here.
"""
- app = FastvideoCausalWan22T2VApplication()
+ pipeline_config: Any = copy.deepcopy(RUNNER_WAN22_T2V_14B.pipeline)
+
+ def load_stand_in(_: object) -> FakeT2VPipeline:
+ return FakeT2VPipeline()
+
+ pipeline_config._target = load_stand_in
+ app = FastvideoCausalWan22T2VApplication(pipeline_config=pipeline_config)
app.init(["--prompt", _PROMPT, "--no-compile"])
diff --git a/integrations_v2/t2v_self_forcing/README.md b/integrations_v2/t2v_self_forcing/README.md
index 84afb0ae6..e62676758 100644
--- a/integrations_v2/t2v_self_forcing/README.md
+++ b/integrations_v2/t2v_self_forcing/README.md
@@ -20,6 +20,19 @@ flashdreams-run-v2 t2v-self-forcing --output-path clip.mp4 \
Arguments after `--` go to the application, and
`flashdreams-run-v2 t2v-self-forcing -- --help` lists them.
+To stream one clip to a browser instead of writing an MP4:
+
+```bash
+uv run --project integrations_v2/t2v_self_forcing flashdreams-run-v2 \
+ t2v-self-forcing --mode webrtc
+```
+
+Open the printed URL, enter a prompt, and select **New session**. The server
+cleans up each completed rollout and waits for another prompt without reloading
+the model. Closing or refreshing the page ends the active rollout but not the
+server; the browser can reconnect and request another session. Stop it with
+Ctrl-C.
+
`--total-blocks` is how many autoregressive blocks to generate, and the run ends
when the session has generated them. The first block decodes 9 frames and every
block after it 12, at 16 frames per second, so seven blocks is about four and a
diff --git a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py
index 67b5e0b8e..11ce0c015 100644
--- a/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py
+++ b/integrations_v2/t2v_self_forcing/t2v_self_forcing/tests/test_stand_in_model.py
@@ -9,8 +9,10 @@
``test_real_model.py``.
"""
+import copy
import shutil
from pathlib import Path
+from typing import Any
import pytest
from self_forcing.config import RUNNER_WAN21_T2V_1PT3B
@@ -36,7 +38,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
integration already ships rather than written down again."""
app = SelfForcingT2VApplication(pipeline_config=FakeT2VPipelineConfig())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert (desc.video_width, desc.video_height) == (
RUNNER_WAN21_T2V_1PT3B.pixel_width,
@@ -48,10 +50,14 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
def test_compilation_can_be_turned_off_for_a_run() -> None:
- """Run against the real config rather than a stand-in, since what this
- covers is the override landing where this model keeps the setting. No model
- is loaded to answer it."""
- app = SelfForcingT2VApplication()
+ """Apply the override to the real config while loading a stand-in model."""
+ pipeline_config: Any = copy.deepcopy(RUNNER_WAN21_T2V_1PT3B.pipeline)
+
+ def load_stand_in(_: object) -> FakeT2VPipeline:
+ return FakeT2VPipeline()
+
+ pipeline_config._target = load_stand_in
+ app = SelfForcingT2VApplication(pipeline_config=pipeline_config)
app.init(["--prompt", _PROMPT, "--no-compile"])
diff --git a/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py b/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py
index 26ee7140b..f3a337b81 100644
--- a/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py
+++ b/integrations_v2/t2v_wan21/t2v_wan21/tests/test_stand_in_model.py
@@ -27,7 +27,7 @@ def test_the_model_says_what_it_generates_without_being_told() -> None:
for a model that does not roll out does not carry one."""
app = Wan21T2VApplication(pipeline_config=FakeT2VPipelineConfig())
- desc = app.session_desc()
+ desc = app.default_session_desc()
assert (desc.video_width, desc.video_height) == (
RUNNER_WAN21_T2V_1PT3B_480P.pixel_width,