From 762da54caa5b2bd31e9da1ce47127b07c7860b71 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 06:58:38 +0000 Subject: [PATCH 01/33] Add public demo application contracts Introduce the flashdreams.demo facade for application-facing demo protocols and thin adapters over the existing runtime demo stack. Cover the Phase 1 contract shape with focused CPU tests. --- flashdreams/flashdreams/demo/__init__.py | 30 +++ flashdreams/flashdreams/demo/application.py | 254 ++++++++++++++++++ .../tests/test_demo_application_api.py | 244 +++++++++++++++++ 3 files changed, 528 insertions(+) create mode 100644 flashdreams/flashdreams/demo/__init__.py create mode 100644 flashdreams/flashdreams/demo/application.py create mode 100644 flashdreams/tests/test_demo_application_api.py diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py new file mode 100644 index 000000000..d62729c9f --- /dev/null +++ b/flashdreams/flashdreams/demo/__init__.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public demo application authoring API.""" + +from flashdreams.demo.application import ( + Application, + ApplicationSession, + DemoAdapterApplication, + FrameOutputSink, + IApplication, + IApplicationSession, + IOutputSink, + IOHandler, + InferenceSessionApplicationAdapter, + RuntimeOutputSinkFrameAdapter, +) + +__all__ = [ + "Application", + "ApplicationSession", + "DemoAdapterApplication", + "FrameOutputSink", + "IApplication", + "IApplicationSession", + "IOutputSink", + "IOHandler", + "InferenceSessionApplicationAdapter", + "RuntimeOutputSinkFrameAdapter", +] diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py new file mode 100644 index 000000000..e2f6c15b7 --- /dev/null +++ b/flashdreams/flashdreams/demo/application.py @@ -0,0 +1,254 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public demo application contracts. + +This module is the author-facing facade above ``flashdreams.runtime.demo``. The +runtime package owns execution, providers, sinks, drivers, and worker affinity; +these protocols name the smaller surface demo applications should eventually +implement directly. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime.demo.outputs import ( + OutputDecision, + OutputSink, + SessionInfo, +) +from flashdreams.runtime.demo.session_inputs import UserInputWindow +from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec, PreparedScenario +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.inputs import InferenceInput +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import ( + StepRequest, + StepRequirements, + StepResult, + step_requirements_from_request, +) + + +@runtime_checkable +class ApplicationSession(Protocol): + """One model session exposed through the public demo API.""" + + def init(self) -> None: + """Initialize per-session resources before the first step.""" + ... + + def session_info(self) -> SessionInfo: + """Return output-facing metadata known after session setup.""" + ... + + def next_step_requirements(self) -> StepRequirements | None: + """Return requirements for the next step, or ``None`` when complete.""" + ... + + def step(self, model_input: InferenceInput) -> StepResult: + """Run one inference step from model-facing inputs.""" + ... + + def reset(self, model_input: InferenceInput | None = None) -> None: + """Reset rollout state when the backend supports it.""" + ... + + def close(self) -> None: + """Release per-session resources.""" + ... + + +@runtime_checkable +class Application(Protocol): + """Public demo application facade.""" + + def init(self, launch_args: Sequence[str]) -> None: + """Initialize application-level launch state.""" + ... + + def create_session(self) -> ApplicationSession: + """Create one model session. + + Public runners must call this on the same worker that will execute the + session so model construction and stepping share worker affinity. + """ + ... + + +@runtime_checkable +class IOHandler(Protocol): + """Public facade over runtime input, output, and transport edges.""" + + def open(self, session_info: SessionInfo) -> None: + """Prepare input/output resources after session setup.""" + ... + + def next_window(self, requirements: StepRequirements) -> UserInputWindow: + """Return the user-input window selected for one model step. + + Model-specific conversion from this window to ``InferenceInput`` remains + the responsibility of ``ModelInputProvider`` in the runtime layer. + """ + ... + + def get_user_input_state(self, modality: str, name: str) -> Any: + """Return the current named input state for interactive applications.""" + ... + + def begin_generation(self, generation: int) -> None: + """Start a new output generation.""" + ... + + def emit_chunk(self, result: StepResult) -> OutputDecision: + """Deliver one generated step result.""" + ... + + def should_exit(self) -> bool: + """Return whether the surrounding run should stop.""" + ... + + def close(self) -> Sequence[OutputArtifact]: + """Finalize resources and return produced artifacts.""" + ... + + +@runtime_checkable +class FrameOutputSink(Protocol): + """Narrow file/comparison tail used by higher-level output handlers.""" + + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + """Consume one timestamped generated chunk.""" + ... + + +@dataclass(slots=True) +class InferenceSessionApplicationAdapter: + """Adapt an existing runtime session to ``ApplicationSession``.""" + + session: InferenceSession + + def init(self) -> None: + init = getattr(self.session, "init", None) + if callable(init): + init() + + def session_info(self) -> SessionInfo: + session_info = getattr(self.session, "session_info", None) + if not callable(session_info): + return SessionInfo() + value = session_info() + if not isinstance(value, SessionInfo): + raise TypeError( + "session.session_info() must return SessionInfo, " + f"got {type(value).__name__}." + ) + return value + + def next_step_requirements(self) -> StepRequirements | None: + next_requirements = getattr(self.session, "next_step_requirements", None) + if callable(next_requirements): + value = next_requirements() + if value is None or isinstance(value, StepRequirements): + return value + raise TypeError( + "session.next_step_requirements() must return StepRequirements " + f"or None, got {type(value).__name__}." + ) + + request = self.session.next_step_request() + if request is None: + return None + if not isinstance(request, StepRequest): + raise TypeError( + "session.next_step_request() must return StepRequest or None, " + f"got {type(request).__name__}." + ) + return step_requirements_from_request( + request, + allow_user_input_window=True, + ) + + def step(self, model_input: InferenceInput) -> StepResult: + return self.session.step(model_input) + + def reset(self, model_input: InferenceInput | None = None) -> None: + self.session.reset(model_input) + + def close(self) -> None: + self.session.close() + + +@dataclass(slots=True) +class DemoAdapterApplication: + """Adapt an existing ``DemoAdapter`` to the public ``Application`` shape.""" + + adapter: DemoAdapter + spec: DemoSpec + _scenario: PreparedScenario | None = field(default=None, init=False, repr=False) + _runtime: InferenceRuntime | None = field(default=None, init=False, repr=False) + + def init(self, launch_args: Sequence[str]) -> None: + del launch_args + config = _require_config(self.spec) + self.adapter.validate_config(config) + self._scenario = self.adapter.prepare_scenario(self.spec) + + def create_session(self) -> ApplicationSession: + scenario = self._scenario + if scenario is None: + self.init(()) + scenario = self._scenario + if scenario is None: + raise RuntimeError("DemoAdapterApplication failed to prepare a scenario.") + runtime = self.adapter.create_runtime(_require_config(self.spec)) + self._runtime = runtime + return InferenceSessionApplicationAdapter( + runtime.start_session(scenario.initial_inputs) + ) + + def close(self) -> None: + if self._runtime is not None: + self._runtime.close() + self._runtime = None + + +@dataclass(slots=True) +class RuntimeOutputSinkFrameAdapter: + """Adapt a runtime ``OutputSink`` to the narrow frame-output tail.""" + + output_sink: OutputSink + + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.output_sink.write(chunk) + + +def _require_config(spec: DemoSpec) -> InferenceConfig: + config = spec.config + if config is None: + raise RuntimeError("DemoSpec.config must be populated before use.") + return config + + +IApplication = Application +IApplicationSession = ApplicationSession +IOutputSink = FrameOutputSink + + +__all__ = [ + "Application", + "ApplicationSession", + "DemoAdapterApplication", + "FrameOutputSink", + "IApplication", + "IApplicationSession", + "IOutputSink", + "IOHandler", + "InferenceSessionApplicationAdapter", + "RuntimeOutputSinkFrameAdapter", +] diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py new file mode 100644 index 000000000..e8d251f67 --- /dev/null +++ b/flashdreams/tests/test_demo_application_api.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import pytest + +from flashdreams.demo import ( + Application, + ApplicationSession, + DemoAdapterApplication, + FrameOutputSink, + IOHandler, + InferenceSessionApplicationAdapter, + RuntimeOutputSinkFrameAdapter, +) +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputField, + InputMapping, + StepRequest, + StepResult, + TimeWindow, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + DemoSpec, + NullOutputSink, + NullOutputSpec, + OutputDecision, + PreparedScenario, + SessionInfo, + UserInputWindow, +) +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import StepRequirements + +pytestmark = pytest.mark.ci_cpu + + +def test_public_demo_contracts_are_importable() -> None: + assert Application.__name__ == "Application" + assert ApplicationSession.__name__ == "ApplicationSession" + assert IOHandler.__name__ == "IOHandler" + assert FrameOutputSink.__name__ == "FrameOutputSink" + + +def test_inference_session_adapter_satisfies_application_session() -> None: + session = _FakeSession() + adapter = InferenceSessionApplicationAdapter(session) + + assert isinstance(adapter, ApplicationSession) + adapter.init() + requirements = adapter.next_step_requirements() + + assert requirements == StepRequirements( + step_index=0, + inference_input_schema=session.inference_input_schema, + ) + result = adapter.step(InferenceInput(step={"chunk_index": 0})) + assert result.step_index == 0 + assert adapter.session_info() == SessionInfo(output_layout="thwc") + adapter.reset() + adapter.close() + assert session.initialized + assert session.reset_called + assert session.closed + + +def test_demo_adapter_application_satisfies_application() -> None: + demo = DemoAdapterApplication( + adapter=_FakeDemoAdapter(), + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ), + ) + + assert isinstance(demo, Application) + demo.init(()) + session = demo.create_session() + + assert isinstance(session, ApplicationSession) + assert session.next_step_requirements() == StepRequirements( + step_index=0, + inference_input_schema=_FakeSession.inference_input_schema, + ) + demo.close() + + +def test_io_handler_protocol_keeps_input_conversion_outside_io() -> None: + handler = _FakeIOHandler() + + assert isinstance(handler, IOHandler) + window = handler.next_window(StepRequirements(step_index=3)) + + assert window.start_s == 3.0 + assert window.end_s == 4.0 + assert handler.get_user_input_state("keyboard", "key_w") is False + + +def test_runtime_output_sink_frame_adapter_satisfies_frame_output_sink() -> None: + output = NullOutputSink(store_results=True) + output.open(SessionInfo()) + adapter = RuntimeOutputSinkFrameAdapter(output) + + assert isinstance(adapter, FrameOutputSink) + adapter.handle_output( + 0.0, + StepResult( + step_index=0, + output="chunk", + frame_count=1, + output_window=TimeWindow(start_s=0.0, end_s=1.0), + ), + ) + + assert output.output_count == 1 + + +class _FakeSession: + inference_input_schema = InferenceInputSchema( + step_fields=(InputField(name="chunk_index"),) + ) + + def __init__(self) -> None: + self.initialized = False + self.closed = False + self.reset_called = False + + def init(self) -> None: + self.initialized = True + + def session_info(self) -> SessionInfo: + return SessionInfo(output_layout="thwc") + + def next_step_request(self) -> StepRequest | None: + return StepRequest( + step_index=0, + inference_input_schema=self.inference_input_schema, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + self.inference_input_schema.require_step(inputs) + return StepResult( + step_index=0, + output="chunk", + frame_count=1, + output_window=TimeWindow(start_s=0.0, end_s=1.0), + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.reset_called = True + + def close(self) -> None: + self.closed = True + + +class _FakeRuntime: + def __init__(self) -> None: + self.session = _FakeSession() + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + assert inputs.global_conditioning["prompt"] == "demo" + return self.session + + def close(self) -> None: + self.closed = True + + +class _FakeDemoAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + canonical_input_schema = CanonicalInputSchema() + + def __init__(self) -> None: + self.runtime = _FakeRuntime() + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + assert config.model_id == self.model_id + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return self.runtime + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + assert spec.model_id == self.model_id + return PreparedScenario( + initial_inputs=InferenceInput(global_conditioning={"prompt": "demo"}), + user_inputs=UserInputs(), + source_schema=UserInputSchema(), + ) + + +class _FakeIOHandler: + def open(self, session_info: SessionInfo) -> None: + del session_info + + def next_window(self, requirements: StepRequirements) -> UserInputWindow: + start_s = float(requirements.step_index) + return UserInputWindow(start_s=start_s, end_s=start_s + 1.0) + + def get_user_input_state(self, modality: str, name: str) -> Any: + assert modality == "keyboard" + assert name == "key_w" + return False + + def begin_generation(self, generation: int) -> None: + assert generation >= 0 + + def emit_chunk(self, result: StepResult) -> OutputDecision: + assert result.step_index >= 0 + return OutputDecision() + + def should_exit(self) -> bool: + return False + + def close(self) -> Sequence[OutputArtifact]: + return () From 991b66e5ca22655aa112b91dda4e6275f80fdbb4 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 06:59:41 +0000 Subject: [PATCH 02/33] Fix linter issues --- flashdreams/flashdreams/demo/__init__.py | 4 ++-- flashdreams/flashdreams/demo/application.py | 2 +- flashdreams/tests/test_demo_application_api.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index d62729c9f..00f1cfb76 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -10,9 +10,9 @@ FrameOutputSink, IApplication, IApplicationSession, - IOutputSink, - IOHandler, InferenceSessionApplicationAdapter, + IOHandler, + IOutputSink, RuntimeOutputSinkFrameAdapter, ) diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py index e2f6c15b7..360a25e13 100644 --- a/flashdreams/flashdreams/demo/application.py +++ b/flashdreams/flashdreams/demo/application.py @@ -15,6 +15,7 @@ from dataclasses import dataclass, field from typing import Any, Protocol, runtime_checkable +from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.demo.outputs import ( OutputDecision, OutputSink, @@ -22,7 +23,6 @@ ) from flashdreams.runtime.demo.session_inputs import UserInputWindow from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec, PreparedScenario -from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.inputs import InferenceInput from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession from flashdreams.runtime.output import OutputArtifact diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index e8d251f67..7c5ccc052 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -13,8 +13,8 @@ ApplicationSession, DemoAdapterApplication, FrameOutputSink, - IOHandler, InferenceSessionApplicationAdapter, + IOHandler, RuntimeOutputSinkFrameAdapter, ) from flashdreams.runtime import ( From 72afb61d051a361493f69926b8f115a6d7c4e334 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 07:10:31 +0000 Subject: [PATCH 03/33] Add public demo Runner facade Expose flashdreams.demo.Runner as an application-facing wrapper over the existing shared demo session helpers. Adapt public Application and IOHandler objects into the runtime demo stack while preserving RuntimeHost worker affinity, StepPipeline execution, metrics, output decisions, and cleanup. Cover direct sync and async Runner usage with focused CPU tests. --- flashdreams/flashdreams/demo/__init__.py | 2 + flashdreams/flashdreams/demo/runner.py | 385 ++++++++++++++++++ .../tests/test_demo_application_api.py | 274 +++++++++++++ 3 files changed, 661 insertions(+) create mode 100644 flashdreams/flashdreams/demo/runner.py diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 00f1cfb76..0c85d4d35 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -15,6 +15,7 @@ IOutputSink, RuntimeOutputSinkFrameAdapter, ) +from flashdreams.demo.runner import Runner __all__ = [ "Application", @@ -27,4 +28,5 @@ "IOHandler", "InferenceSessionApplicationAdapter", "RuntimeOutputSinkFrameAdapter", + "Runner", ] diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py new file mode 100644 index 000000000..6426f9c23 --- /dev/null +++ b/flashdreams/flashdreams/demo/runner.py @@ -0,0 +1,385 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public runner facade for demo applications.""" + +from __future__ import annotations + +import asyncio +import inspect +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any, cast + +from flashdreams.runtime.canonical import CanonicalInputSchema +from flashdreams.runtime.config import InferenceConfig +from flashdreams.runtime.demo.drivers import ( + BatchSessionDriver, + run_demo_session, + run_demo_session_async, +) +from flashdreams.runtime.demo.host import ModelWarmupPlan, RuntimeHost +from flashdreams.runtime.demo.outputs import OutputDecision, SessionInfo +from flashdreams.runtime.demo.pipeline import StepPipeline +from flashdreams.runtime.demo.run_modes import ( + AsyncSessionDriver, + InMemorySessionMetricsRecorder, + NoopTransportService, + RunContext, + RunMode, + RunModeCapabilities, + RunResult, + SessionDriver, + SessionEdges, + SessionMetricsRecorder, + SingleSessionAdmissionPolicy, +) +from flashdreams.runtime.demo.session_inputs import ( + PreparedStep, + ProviderCapabilities, + UserInputWindow, +) +from flashdreams.runtime.demo.spec import ( + DemoAdapter, + DemoSpec, + NullOutputSpec, + PreparedScenario, +) +from flashdreams.runtime.inputs import ( + InferenceInput, + InferenceInputSchema, + UserInputSchema, +) +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.mapping import InputMapping +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import StepRequirements, StepResult + +from .application import Application, ApplicationSession, IOHandler + + +@dataclass(slots=True) +class Runner: + """Run a public demo application through the shared demo runtime. + + This facade exists for application authors. It adapts ``Application`` and + ``IOHandler`` to the existing runtime ``run_demo_session`` helpers, so the + model worker boundary, ``StepPipeline``, metrics, output decisions, and + cleanup behavior stay in the runtime implementation. + """ + + io_handler: IOHandler + app: Application + launch_args: Sequence[str] = () + host: RuntimeHost | None = None + metrics: SessionMetricsRecorder | None = None + pipeline: StepPipeline | None = None + run_mode: RunMode | None = None + model_id: str | None = None + + def run(self) -> RunResult: + """Run one session and return its shared runtime result.""" + if _run_mode_is_async(self._selected_run_mode()): + return asyncio.run(self.run_async()) + return self._run_sync() + + async def run_async(self) -> RunResult: + """Run one session through the async helper when the run mode needs it.""" + run_mode = self._selected_run_mode() + if not _run_mode_is_async(run_mode): + return self._run_sync(run_mode=run_mode) + + host, owns_host = self._selected_host() + context = self._create_context(host) + spec = self._create_spec(run_mode) + scenario = _runner_scenario() + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + try: + self.app.init(tuple(self.launch_args)) + return await run_demo_session_async( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=run_mode, + pipeline=self.pipeline or StepPipeline(), + ) + finally: + await context.close_async() + if owns_host: + host.close() + + def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: + selected_run_mode = run_mode or self._selected_run_mode() + host, owns_host = self._selected_host() + context = self._create_context(host) + spec = self._create_spec(selected_run_mode) + scenario = _runner_scenario() + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + try: + self.app.init(tuple(self.launch_args)) + return run_demo_session( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=selected_run_mode, + pipeline=self.pipeline or StepPipeline(), + ) + finally: + context.close() + if owns_host: + host.close() + + def _selected_run_mode(self) -> RunMode: + if self.run_mode is not None: + return self.run_mode + return _IOHandlerRunMode(self.io_handler) + + def _selected_host(self) -> tuple[RuntimeHost, bool]: + if self.host is not None: + return self.host, False + return RuntimeHost(_ApplicationRuntime(self.app)), True + + def _create_context(self, host: RuntimeHost) -> RunContext: + return RunContext( + host=host, + run_metrics=self.metrics or InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy, + ), + model_warmup_plan=ModelWarmupPlan(), + ) + + def _create_spec(self, run_mode: RunMode) -> DemoSpec: + model_id = self.model_id or _application_model_id(self.app) + return DemoSpec( + model_id=model_id, + input_mode=run_mode.name, + output=NullOutputSpec(), + config=InferenceConfig(model_id=model_id), + ) + + +@dataclass(slots=True) +class _ApplicationRuntime: + app: Application + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + session = self.app.create_session() + if not isinstance(session, ApplicationSession): + raise TypeError( + "Application.create_session() must return ApplicationSession, " + f"got {type(session).__name__}." + ) + session.init() + return cast(InferenceSession, session) + + def close(self) -> None: + close = getattr(self.app, "close", None) + if callable(close): + close() + + +@dataclass(slots=True) +class _RunnerDemoAdapter: + app: Application + spec: DemoSpec + scenario: PreparedScenario + + @property + def model_id(self) -> str: + return self.spec.model_id + + @property + def inference_input_schema(self) -> InferenceInputSchema: + return InferenceInputSchema() + + @property + def canonical_input_schema(self) -> CanonicalInputSchema: + return CanonicalInputSchema() + + def default_input_mapping(self) -> InputMapping | None: + return None + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.spec.model_id: + raise ValueError(f"Unsupported model_id={config.model_id!r}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + return cast(InferenceRuntime, _ApplicationRuntime(self.app)) + + def supported_input_modes(self) -> tuple[str, ...]: + return (self.spec.input_mode,) + + def supported_output_modes(self) -> tuple[str, ...]: + return (self.spec.output.mode,) + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + if spec != self.spec: + raise ValueError("Runner received an unexpected DemoSpec.") + return self.scenario + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: Any, + ) -> "_RunnerModelInputProvider": + del spec, scenario + return _RunnerModelInputProvider() + + +@dataclass(slots=True) +class _RunnerModelInputProvider: + capabilities: ProviderCapabilities = field( + default_factory=lambda: ProviderCapabilities( + supports_recorded_input=True, + inference_input_schema=InferenceInputSchema(), + ) + ) + closed: bool = False + + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput() + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + return PreparedStep( + inference_input=InferenceInput( + step={ + "step_index": request.step_index, + "user_window": user_window, + }, + metadata=request.metadata, + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.closed = True + + +@dataclass(slots=True) +class _IOHandlerRunMode: + io_handler: IOHandler + driver: SessionDriver | AsyncSessionDriver = field( + default_factory=BatchSessionDriver + ) + name: str = "public-runner" + capabilities: RunModeCapabilities = field( + default_factory=lambda: RunModeCapabilities(supports_artifacts=True) + ) + + def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None: + del spec, adapter + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: Any, + adapter: DemoAdapter, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: DemoAdapter, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + return RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy, + ), + model_warmup_plan=model_warmup_plan, + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + provider: Any, + adapter: DemoAdapter, + ) -> SessionEdges: + del spec, scenario, provider, adapter + return SessionEdges( + input_source=_IOHandlerBatchInputSource(self.io_handler), + output_sink=_IOHandlerOutputSink(self.io_handler), + cleanup_tasks=context.cleanup_tasks, + transport=NoopTransportService(), + ) + + def select_driver(self) -> SessionDriver | AsyncSessionDriver: + return self.driver + + +@dataclass(slots=True) +class _IOHandlerBatchInputSource: + io_handler: IOHandler + is_finite: bool = False + is_deterministic: bool = False + user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) + + def is_finished(self) -> bool: + return self.io_handler.should_exit() + + def next_window(self, request: StepRequirements) -> UserInputWindow: + return self.io_handler.next_window(request) + + +@dataclass(slots=True) +class _IOHandlerOutputSink: + io_handler: IOHandler + produces_artifacts: bool = True + _generation_started: bool = field(default=False, init=False, repr=False) + + def open(self, session_info: SessionInfo) -> None: + self.io_handler.open(session_info) + + def begin_generation(self, generation: int) -> None: + self._generation_started = True + self.io_handler.begin_generation(generation) + + def write(self, result: StepResult) -> OutputDecision: + if not self._generation_started: + self.begin_generation(0) + return self.io_handler.emit_chunk(result) + + def close(self) -> Sequence[OutputArtifact]: + return self.io_handler.close() + + +def _run_mode_is_async(run_mode: RunMode) -> bool: + driver = run_mode.select_driver() + return inspect.iscoroutinefunction(driver.run_one_session) + + +def _application_model_id(app: Application) -> str: + model_id = getattr(app, "model_id", None) + if isinstance(model_id, str) and model_id.strip(): + return model_id + return app.__class__.__name__ or "demo-application" + + +def _runner_scenario() -> PreparedScenario: + return PreparedScenario(initial_inputs=InferenceInput()) + + +__all__ = ["Runner"] diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 7c5ccc052..4509bb36a 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -3,7 +3,9 @@ from __future__ import annotations +import threading from collections.abc import Sequence +from dataclasses import dataclass, field from typing import Any import pytest @@ -15,6 +17,7 @@ FrameOutputSink, InferenceSessionApplicationAdapter, IOHandler, + Runner, RuntimeOutputSinkFrameAdapter, ) from flashdreams.runtime import ( @@ -34,12 +37,20 @@ UserInputSchema, ) from flashdreams.runtime.demo import ( + BatchSessionDriver, DemoSpec, + InMemorySessionMetricsRecorder, NullOutputSink, NullOutputSpec, OutputDecision, PreparedScenario, + RunContext, + RunModeCapabilities, + RunResult, + SessionEdges, SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, UserInputWindow, ) from flashdreams.runtime.output import OutputArtifact @@ -129,6 +140,54 @@ def test_runtime_output_sink_frame_adapter_satisfies_frame_output_sink() -> None assert output.output_count == 1 +def test_runner_run_drives_public_app_through_shared_runtime_path() -> None: + caller_thread_id = threading.get_ident() + app = _RunnerFakeApplication(total_steps=2) + io_handler = _RecordingIOHandler() + + result = Runner( + io_handler=io_handler, + app=app, + launch_args=("--quality", "fast"), + model_id="fake-runner", + ).run() + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 2 + assert app.launch_args == ("--quality", "fast") + assert app.init_thread_id == caller_thread_id + assert app.session is not None + assert app.session.init_thread_id != caller_thread_id + assert app.session.step_thread_ids == (app.session.init_thread_id,) * 2 + assert app.session.closed + assert io_handler.opened_with == [SessionInfo(output_layout="thwc")] + assert io_handler.requested_steps == [0, 1] + assert io_handler.begin_generations == [0] + assert io_handler.emitted_steps == [0, 1] + assert io_handler.closed + + +@pytest.mark.asyncio +async def test_runner_run_async_delegates_to_async_session_helper() -> None: + app = _RunnerFakeApplication(total_steps=1) + io_handler = _RecordingIOHandler() + run_mode = _AsyncRecordingRunMode(io_handler) + + result = await Runner( + io_handler=io_handler, + app=app, + run_mode=run_mode, + model_id="fake-runner", + ).run_async() + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert run_mode.driver.called + assert io_handler.emitted_steps == [0] + + class _FakeSession: inference_input_schema = InferenceInputSchema( step_fields=(InputField(name="chunk_index"),) @@ -242,3 +301,218 @@ def should_exit(self) -> bool: def close(self) -> Sequence[OutputArtifact]: return () + + +class _RunnerFakeApplication: + model_id = "runner-fake" + + def __init__(self, *, total_steps: int) -> None: + self.total_steps = total_steps + self.launch_args: tuple[str, ...] = () + self.init_thread_id: int | None = None + self.session: _RunnerFakeSession | None = None + + def init(self, launch_args: Sequence[str]) -> None: + self.launch_args = tuple(launch_args) + self.init_thread_id = threading.get_ident() + + def create_session(self) -> "_RunnerFakeSession": + self.session = _RunnerFakeSession(total_steps=self.total_steps) + return self.session + + +class _RunnerFakeSession: + def __init__(self, *, total_steps: int) -> None: + self.total_steps = total_steps + self.step_index = 0 + self.init_thread_id: int | None = None + self.step_thread_ids: tuple[int, ...] = () + self.closed = False + + def init(self) -> None: + self.init_thread_id = threading.get_ident() + + def session_info(self) -> SessionInfo: + return SessionInfo(output_layout="thwc") + + def next_step_requirements(self) -> StepRequirements | None: + if self.step_index >= self.total_steps: + return None + return StepRequirements(step_index=self.step_index) + + def step(self, model_input: InferenceInput) -> StepResult: + assert model_input.step["step_index"] == self.step_index + assert isinstance(model_input.step["user_window"], UserInputWindow) + self.step_thread_ids = (*self.step_thread_ids, threading.get_ident()) + result = StepResult( + step_index=self.step_index, + output=f"runner-chunk-{self.step_index}", + frame_count=1, + output_window=TimeWindow( + start_s=float(self.step_index), + end_s=float(self.step_index + 1), + ), + ) + self.step_index += 1 + return result + + def reset(self, model_input: InferenceInput | None = None) -> None: + del model_input + self.step_index = 0 + + def close(self) -> None: + self.closed = True + + +class _RecordingIOHandler: + def __init__(self) -> None: + self.opened_with: list[SessionInfo] = [] + self.requested_steps: list[int] = [] + self.begin_generations: list[int] = [] + self.emitted_steps: list[int] = [] + self.closed = False + + def open(self, session_info: SessionInfo) -> None: + self.opened_with.append(session_info) + + def next_window(self, requirements: StepRequirements) -> UserInputWindow: + self.requested_steps.append(requirements.step_index) + start_s = float(requirements.step_index) + return UserInputWindow(start_s=start_s, end_s=start_s + 1.0) + + def get_user_input_state(self, modality: str, name: str) -> Any: + del modality, name + return None + + def begin_generation(self, generation: int) -> None: + self.begin_generations.append(generation) + + def emit_chunk(self, result: StepResult) -> OutputDecision: + self.emitted_steps.append(result.step_index) + return OutputDecision() + + def should_exit(self) -> bool: + return False + + def close(self) -> Sequence[OutputArtifact]: + self.closed = True + return () + + +@dataclass(slots=True) +class _AsyncRecordingRunMode: + io_handler: _RecordingIOHandler + name: str = "async-public-runner" + capabilities: RunModeCapabilities = field( + default_factory=lambda: RunModeCapabilities(supports_artifacts=True) + ) + driver: "_AsyncBatchDriver" = field(default_factory=lambda: _AsyncBatchDriver()) + + def validate_run(self, *, spec: DemoSpec, adapter: Any) -> None: + del spec, adapter + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: Any, + adapter: Any, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: Any, + host: Any, + model_warmup_plan: Any, + ) -> RunContext: + del spec, adapter + return RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy, + ), + model_warmup_plan=model_warmup_plan, + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + provider: Any, + adapter: Any, + ) -> SessionEdges: + del spec, scenario, provider, adapter + return SessionEdges( + input_source=_AsyncRunModeInputSource(self.io_handler), + output_sink=_AsyncRunModeOutputSink(self.io_handler), + cleanup_tasks=context.cleanup_tasks, + ) + + def select_driver(self) -> "_AsyncBatchDriver": + return self.driver + + +class _AsyncBatchDriver: + def __init__(self) -> None: + self.called = False + + async def run_one_session( + self, + *, + host: Any, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + self.called = True + return BatchSessionDriver().run_one_session( + host=host, + provider=provider, + session_edges=session_edges, + pipeline=pipeline, + ) + + +class _AsyncRunModeInputSource: + is_finite = False + is_deterministic = False + user_input_schema = UserInputSchema() + + def __init__(self, io_handler: _RecordingIOHandler) -> None: + self._io_handler = io_handler + + def is_finished(self) -> bool: + return self._io_handler.should_exit() + + def next_window(self, request: StepRequirements) -> UserInputWindow: + return self._io_handler.next_window(request) + + +class _AsyncRunModeOutputSink: + produces_artifacts = True + + def __init__(self, io_handler: _RecordingIOHandler) -> None: + self._io_handler = io_handler + self._begun = False + + def open(self, session_info: SessionInfo) -> None: + self._io_handler.open(session_info) + + def begin_generation(self, generation: int) -> None: + self._begun = True + self._io_handler.begin_generation(generation) + + def write(self, result: StepResult) -> OutputDecision: + if not self._begun: + self.begin_generation(0) + return self._io_handler.emit_chunk(result) + + def close(self) -> Sequence[OutputArtifact]: + return self._io_handler.close() From 04a450d73f27ae8d03538aeffd4ee59cb842b866 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 07:15:05 +0000 Subject: [PATCH 04/33] Fix public demo application lifecycle Add application-level cleanup to the public demo contract and make DemoAdapterApplication close every runtime it creates. Reject unsupported launch args instead of silently dropping them, and cover repeated session creation and cleanup with focused CPU tests. --- flashdreams/flashdreams/demo/application.py | 32 ++++++++++--- flashdreams/flashdreams/demo/runner.py | 4 +- .../tests/test_demo_application_api.py | 47 ++++++++++++++++++- 3 files changed, 73 insertions(+), 10 deletions(-) diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py index 360a25e13..8441e9d78 100644 --- a/flashdreams/flashdreams/demo/application.py +++ b/flashdreams/flashdreams/demo/application.py @@ -79,6 +79,10 @@ def create_session(self) -> ApplicationSession: """ ... + def close(self) -> None: + """Release application-level resources after all sessions have closed.""" + ... + @runtime_checkable class IOHandler(Protocol): @@ -190,10 +194,18 @@ class DemoAdapterApplication: adapter: DemoAdapter spec: DemoSpec _scenario: PreparedScenario | None = field(default=None, init=False, repr=False) - _runtime: InferenceRuntime | None = field(default=None, init=False, repr=False) + _runtimes: list[InferenceRuntime] = field( + default_factory=list, + init=False, + repr=False, + ) def init(self, launch_args: Sequence[str]) -> None: - del launch_args + if launch_args: + raise ValueError( + "DemoAdapterApplication does not support launch arguments; " + "configure the DemoSpec before constructing the application." + ) config = _require_config(self.spec) self.adapter.validate_config(config) self._scenario = self.adapter.prepare_scenario(self.spec) @@ -206,15 +218,23 @@ def create_session(self) -> ApplicationSession: if scenario is None: raise RuntimeError("DemoAdapterApplication failed to prepare a scenario.") runtime = self.adapter.create_runtime(_require_config(self.spec)) - self._runtime = runtime + self._runtimes.append(runtime) return InferenceSessionApplicationAdapter( runtime.start_session(scenario.initial_inputs) ) def close(self) -> None: - if self._runtime is not None: - self._runtime.close() - self._runtime = None + errors: list[Exception] = [] + while self._runtimes: + runtime = self._runtimes.pop() + try: + runtime.close() + except Exception as exc: + errors.append(exc) + if errors: + raise RuntimeError( + f"DemoAdapterApplication.close failed for {len(errors)} runtime(s)." + ) from errors[0] @dataclass(slots=True) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 6426f9c23..63ffd7adf 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -177,9 +177,7 @@ def start_session(self, inputs: InferenceInput) -> InferenceSession: return cast(InferenceSession, session) def close(self) -> None: - close = getattr(self.app, "close", None) - if callable(close): - close() + self.app.close() @dataclass(slots=True) diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 4509bb36a..571b9f1d5 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -89,8 +89,9 @@ def test_inference_session_adapter_satisfies_application_session() -> None: def test_demo_adapter_application_satisfies_application() -> None: + adapter = _FakeDemoAdapter() demo = DemoAdapterApplication( - adapter=_FakeDemoAdapter(), + adapter=adapter, spec=DemoSpec( model_id="fake-demo", input_mode="replay", @@ -108,6 +109,42 @@ def test_demo_adapter_application_satisfies_application() -> None: inference_input_schema=_FakeSession.inference_input_schema, ) demo.close() + assert adapter.runtime.closed + + +def test_demo_adapter_application_closes_all_created_runtimes() -> None: + adapter = _FakeDemoAdapter() + demo = DemoAdapterApplication( + adapter=adapter, + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ), + ) + + demo.init(()) + demo.create_session() + demo.create_session() + demo.close() + demo.close() + + assert len(adapter.runtimes) == 2 + assert [runtime.closed for runtime in adapter.runtimes] == [True, True] + + +def test_demo_adapter_application_rejects_unsupported_launch_args() -> None: + demo = DemoAdapterApplication( + adapter=_FakeDemoAdapter(), + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ), + ) + + with pytest.raises(ValueError, match="does not support launch arguments"): + demo.init(("--device", "cuda:0")) def test_io_handler_protocol_keeps_input_conversion_outside_io() -> None: @@ -161,6 +198,7 @@ def test_runner_run_drives_public_app_through_shared_runtime_path() -> None: assert app.session.init_thread_id != caller_thread_id assert app.session.step_thread_ids == (app.session.init_thread_id,) * 2 assert app.session.closed + assert app.closed assert io_handler.opened_with == [SessionInfo(output_layout="thwc")] assert io_handler.requested_steps == [0, 1] assert io_handler.begin_generations == [0] @@ -250,6 +288,7 @@ class _FakeDemoAdapter: def __init__(self) -> None: self.runtime = _FakeRuntime() + self.runtimes: list[_FakeRuntime] = [] def supported_input_modes(self) -> tuple[str, ...]: return ("replay",) @@ -265,6 +304,8 @@ def validate_config(self, config: InferenceConfig) -> None: def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: self.validate_config(config) + self.runtime = _FakeRuntime() + self.runtimes.append(self.runtime) return self.runtime def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: @@ -311,6 +352,7 @@ def __init__(self, *, total_steps: int) -> None: self.launch_args: tuple[str, ...] = () self.init_thread_id: int | None = None self.session: _RunnerFakeSession | None = None + self.closed = False def init(self, launch_args: Sequence[str]) -> None: self.launch_args = tuple(launch_args) @@ -320,6 +362,9 @@ def create_session(self) -> "_RunnerFakeSession": self.session = _RunnerFakeSession(total_steps=self.total_steps) return self.session + def close(self) -> None: + self.closed = True + class _RunnerFakeSession: def __init__(self, *, total_steps: int) -> None: From cd7f35d68ea2fd0f82813adaf64d058c38bb3e47 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 07:25:01 +0000 Subject: [PATCH 05/33] Add public demo IO handler factories Introduce public replay, native-window, and WebRTC IO factory shapes under flashdreams.demo. Keep WebRTC server-shaped so handlers are produced per connection, and move the Runner IO run-mode adapter into the shared factory module. Cover factory-produced replay handlers, metric forwarding, native factory shape, and WebRTC server callback behavior with focused CPU tests. --- flashdreams/flashdreams/demo/__init__.py | 16 + flashdreams/flashdreams/demo/io.py | 372 ++++++++++++++++++ flashdreams/flashdreams/demo/runner.py | 123 +----- .../tests/test_demo_application_api.py | 84 ++++ 4 files changed, 479 insertions(+), 116 deletions(-) create mode 100644 flashdreams/flashdreams/demo/io.py diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 0c85d4d35..798038f8f 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -15,6 +15,15 @@ IOutputSink, RuntimeOutputSinkFrameAdapter, ) +from flashdreams.demo.io import ( + IOHandlerServer, + NativeWindowIOHandler, + ReplayIOHandler, + WebRTCIOHandlerServer, + create_native_window_io_handler, + create_replay_io_handler, + create_webrtc_io_handler, +) from flashdreams.demo.runner import Runner __all__ = [ @@ -26,7 +35,14 @@ "IApplicationSession", "IOutputSink", "IOHandler", + "IOHandlerServer", "InferenceSessionApplicationAdapter", + "NativeWindowIOHandler", + "ReplayIOHandler", "RuntimeOutputSinkFrameAdapter", "Runner", + "WebRTCIOHandlerServer", + "create_native_window_io_handler", + "create_replay_io_handler", + "create_webrtc_io_handler", ] diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py new file mode 100644 index 000000000..181cffada --- /dev/null +++ b/flashdreams/flashdreams/demo/io.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public demo IO handler factories.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from flashdreams.runtime.demo.drivers import BatchSessionDriver +from flashdreams.runtime.demo.host import ModelWarmupPlan, RuntimeHost +from flashdreams.runtime.demo.outputs import ( + NullOutputSink, + OutputDecision, + OutputSink, + SessionInfo, +) +from flashdreams.runtime.demo.run_modes import ( + AsyncSessionDriver, + InMemorySessionMetricsRecorder, + NoopTransportService, + RunContext, + RunModeCapabilities, + RunResult, + SessionDriver, + SessionEdges, + SingleSessionAdmissionPolicy, +) +from flashdreams.runtime.demo.session_inputs import UserInputWindow +from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec +from flashdreams.runtime.inputs import UserInputs, UserInputSchema +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import StepRequirements, StepResult + +from .application import FrameOutputSink, IOHandler + +RunSessionCallback = Callable[[IOHandler], RunResult] + + +@runtime_checkable +class IOHandlerServer(Protocol): + """Server-shaped IO factory for transports that create handlers per peer.""" + + def serve(self, run_session: RunSessionCallback) -> RunResult: + """Serve one or more sessions by passing handlers to ``run_session``.""" + ... + + +@dataclass(slots=True) +class ReplayIOHandler: + """Batch IO handler for replay/null-style public runner sessions.""" + + replay_log: UserInputs | None = None + output_sink: OutputSink | FrameOutputSink | None = None + metric_output_sink: FrameOutputSink | None = None + is_finite: bool = True + is_deterministic: bool = True + user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) + _input_source: "_ReplayIOInputSource" = field(init=False, repr=False) + _output_sink: OutputSink | FrameOutputSink = field(init=False, repr=False) + _opened_session_info: SessionInfo | None = field( + default=None, + init=False, + repr=False, + ) + _generation: int | None = field(default=None, init=False, repr=False) + _closed: bool = field(default=False, init=False, repr=False) + + def __post_init__(self) -> None: + self._input_source = _ReplayIOInputSource(self.replay_log or UserInputs()) + self._output_sink = self.output_sink or NullOutputSink() + + @property + def run_mode(self) -> "IOHandlerRunMode": + """Return the runtime run mode backing this IO handler.""" + return IOHandlerRunMode( + io_handler=self, + name="replay", + capabilities=RunModeCapabilities( + requires_finite_input=True, + supports_artifacts=True, + ), + ) + + def open(self, session_info: SessionInfo) -> None: + self._opened_session_info = session_info + open_output = getattr(self._output_sink, "open", None) + if callable(open_output): + open_output(session_info) + + def next_window(self, requirements: StepRequirements) -> UserInputWindow: + return self._input_source.next_window(requirements) + + def get_user_input_state(self, modality: str, name: str) -> Any: + del modality, name + return None + + def begin_generation(self, generation: int) -> None: + self._generation = generation + begin_generation = getattr(self._output_sink, "begin_generation", None) + if callable(begin_generation): + begin_generation(generation) + + def emit_chunk(self, result: StepResult) -> OutputDecision: + write = getattr(self._output_sink, "write", None) + if callable(write): + decision = write(result) + if not isinstance(decision, OutputDecision): + raise TypeError( + "OutputSink.write must return OutputDecision, " + f"got {type(decision).__name__}." + ) + else: + timestamp_s = _result_timestamp_s(result) + handle_output = getattr(self._output_sink, "handle_output") + handle_output(timestamp_s, result) + decision = OutputDecision() + if self.metric_output_sink is not None: + self.metric_output_sink.handle_output(_result_timestamp_s(result), result) + return decision + + def should_exit(self) -> bool: + return self._closed + + def close(self) -> Sequence[OutputArtifact]: + self._closed = True + close = getattr(self._output_sink, "close", None) + if callable(close): + artifacts = close() + if artifacts is None: + return () + return tuple(artifacts) + return () + + +@dataclass(slots=True) +class NativeWindowIOHandler(ReplayIOHandler): + """Placeholder native-window factory result until native edges are adopted.""" + + @property + def run_mode(self) -> "IOHandlerRunMode": + return IOHandlerRunMode( + io_handler=self, + name="native-window", + capabilities=RunModeCapabilities(supports_artifacts=True), + ) + + +@dataclass(slots=True) +class WebRTCIOHandlerServer: + """Server-shaped testable facade for WebRTC-style per-connection handlers.""" + + host: str + port: int + viewport_size: tuple[int, int] + handlers: Sequence[IOHandler] = () + + def __post_init__(self) -> None: + if not self.host.strip(): + raise ValueError("WebRTC host must be non-empty.") + if not (0 < int(self.port) < 65536): + raise ValueError("WebRTC port must be between 1 and 65535.") + width, height = self.viewport_size + if width <= 0 or height <= 0: + raise ValueError("WebRTC viewport dimensions must be > 0.") + self.handlers = tuple(self.handlers) + + def serve(self, run_session: RunSessionCallback) -> RunResult: + result = RunResult(status="completed") + for handler in self.handlers: + result = run_session(handler) + if result.status not in {"completed", "skipped"}: + return result + return result + + +@dataclass(slots=True) +class IOHandlerRunMode: + """Runtime run-mode adapter for public ``IOHandler`` instances.""" + + io_handler: IOHandler + name: str = "public-runner" + capabilities: RunModeCapabilities = field( + default_factory=lambda: RunModeCapabilities(supports_artifacts=True) + ) + driver: SessionDriver | AsyncSessionDriver = field( + default_factory=BatchSessionDriver + ) + + def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None: + del spec, adapter + + def validate_session( + self, + *, + spec: DemoSpec, + scenario: Any, + adapter: DemoAdapter, + provider: Any, + ) -> None: + del spec, scenario, adapter, provider + + def create_run_context( + self, + *, + spec: DemoSpec, + adapter: DemoAdapter, + host: RuntimeHost, + model_warmup_plan: ModelWarmupPlan, + ) -> RunContext: + del spec, adapter + return RunContext( + host=host, + run_metrics=InMemorySessionMetricsRecorder(), + admission=SingleSessionAdmissionPolicy( + health_check=lambda: host.is_healthy, + ), + model_warmup_plan=model_warmup_plan, + ) + + def create_session_edges( + self, + *, + context: RunContext, + spec: DemoSpec, + scenario: Any, + provider: Any, + adapter: DemoAdapter, + ) -> SessionEdges: + del spec, scenario, provider, adapter + return SessionEdges( + input_source=IOHandlerBatchInputSource(self.io_handler), + output_sink=IOHandlerOutputSink(self.io_handler), + cleanup_tasks=context.cleanup_tasks, + transport=NoopTransportService(), + ) + + def select_driver(self) -> SessionDriver | AsyncSessionDriver: + return self.driver + + +@dataclass(slots=True) +class IOHandlerBatchInputSource: + """Batch input source adapter for public ``IOHandler`` windows.""" + + io_handler: IOHandler + + @property + def is_finite(self) -> bool: + return bool(getattr(self.io_handler, "is_finite", False)) + + @property + def is_deterministic(self) -> bool: + return bool(getattr(self.io_handler, "is_deterministic", False)) + + @property + def user_input_schema(self) -> UserInputSchema: + schema = getattr(self.io_handler, "user_input_schema", None) + if isinstance(schema, UserInputSchema): + return schema + return UserInputSchema() + + def is_finished(self) -> bool: + return self.io_handler.should_exit() + + def next_window(self, request: StepRequirements) -> UserInputWindow: + return self.io_handler.next_window(request) + + +@dataclass(slots=True) +class IOHandlerOutputSink: + """Output sink adapter for public ``IOHandler`` chunks.""" + + io_handler: IOHandler + produces_artifacts: bool = True + _generation_started: bool = field(default=False, init=False, repr=False) + + def open(self, session_info: SessionInfo) -> None: + self.io_handler.open(session_info) + + def begin_generation(self, generation: int) -> None: + self._generation_started = True + self.io_handler.begin_generation(generation) + + def write(self, result: StepResult) -> OutputDecision: + if not self._generation_started: + self.begin_generation(0) + return self.io_handler.emit_chunk(result) + + def close(self) -> Sequence[OutputArtifact]: + return self.io_handler.close() + + +def create_replay_io_handler( + replay_log: UserInputs | None = None, + output_sink: OutputSink | FrameOutputSink | None = None, + metric_output_sink: FrameOutputSink | None = None, +) -> ReplayIOHandler: + """Create a public replay IO handler.""" + return ReplayIOHandler( + replay_log=replay_log, + output_sink=output_sink, + metric_output_sink=metric_output_sink, + ) + + +def create_native_window_io_handler( + viewport_size: tuple[int, int], +) -> NativeWindowIOHandler: + """Create a placeholder native-window-shaped IO handler. + + The real native-window runtime wiring is a later migration phase. Keeping the + factory name public now lets callers select the mode without importing + runtime internals. + """ + width, height = viewport_size + if width <= 0 or height <= 0: + raise ValueError("Native window viewport dimensions must be > 0.") + return NativeWindowIOHandler() + + +def create_webrtc_io_handler( + host: str, + port: int, + viewport_size: tuple[int, int], + *, + handlers: Sequence[IOHandler] = (), +) -> IOHandlerServer: + """Create a server-shaped WebRTC IO facade.""" + return WebRTCIOHandlerServer( + host=host, + port=port, + viewport_size=viewport_size, + handlers=handlers, + ) + + +class _ReplayIOInputSource: + def __init__(self, replay_log: UserInputs) -> None: + self._replay_log = replay_log + + def next_window(self, requirements: StepRequirements) -> UserInputWindow: + start_s = float(requirements.step_index) + end_s = float(requirements.step_index + requirements.input_frame_count) + return UserInputWindow( + start_s=start_s, + end_s=end_s, + inputs=self._replay_log, + ) + + +def _result_timestamp_s(result: StepResult) -> float: + if result.output_window is None: + return float(result.step_index) + return result.output_window.start_s + + +__all__ = [ + "IOHandlerBatchInputSource", + "IOHandlerOutputSink", + "IOHandlerRunMode", + "IOHandlerServer", + "NativeWindowIOHandler", + "ReplayIOHandler", + "RunSessionCallback", + "WebRTCIOHandlerServer", + "create_native_window_io_handler", + "create_replay_io_handler", + "create_webrtc_io_handler", +] diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 63ffd7adf..cd898da33 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -14,23 +14,16 @@ from flashdreams.runtime.canonical import CanonicalInputSchema from flashdreams.runtime.config import InferenceConfig from flashdreams.runtime.demo.drivers import ( - BatchSessionDriver, run_demo_session, run_demo_session_async, ) from flashdreams.runtime.demo.host import ModelWarmupPlan, RuntimeHost -from flashdreams.runtime.demo.outputs import OutputDecision, SessionInfo from flashdreams.runtime.demo.pipeline import StepPipeline from flashdreams.runtime.demo.run_modes import ( - AsyncSessionDriver, InMemorySessionMetricsRecorder, - NoopTransportService, RunContext, RunMode, - RunModeCapabilities, RunResult, - SessionDriver, - SessionEdges, SessionMetricsRecorder, SingleSessionAdmissionPolicy, ) @@ -39,23 +32,17 @@ ProviderCapabilities, UserInputWindow, ) -from flashdreams.runtime.demo.spec import ( - DemoAdapter, - DemoSpec, - NullOutputSpec, - PreparedScenario, -) +from flashdreams.runtime.demo.spec import DemoSpec, NullOutputSpec, PreparedScenario from flashdreams.runtime.inputs import ( InferenceInput, InferenceInputSchema, - UserInputSchema, ) from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession from flashdreams.runtime.mapping import InputMapping -from flashdreams.runtime.output import OutputArtifact -from flashdreams.runtime.types import StepRequirements, StepResult +from flashdreams.runtime.types import StepRequirements from .application import Application, ApplicationSession, IOHandler +from .io import IOHandlerRunMode @dataclass(slots=True) @@ -134,7 +121,10 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: def _selected_run_mode(self) -> RunMode: if self.run_mode is not None: return self.run_mode - return _IOHandlerRunMode(self.io_handler) + run_mode = getattr(self.io_handler, "run_mode", None) + if run_mode is not None: + return cast(RunMode, run_mode) + return IOHandlerRunMode(self.io_handler) def _selected_host(self) -> tuple[RuntimeHost, bool]: if self.host is not None: @@ -265,105 +255,6 @@ def close(self) -> None: self.closed = True -@dataclass(slots=True) -class _IOHandlerRunMode: - io_handler: IOHandler - driver: SessionDriver | AsyncSessionDriver = field( - default_factory=BatchSessionDriver - ) - name: str = "public-runner" - capabilities: RunModeCapabilities = field( - default_factory=lambda: RunModeCapabilities(supports_artifacts=True) - ) - - def validate_run(self, *, spec: DemoSpec, adapter: DemoAdapter) -> None: - del spec, adapter - - def validate_session( - self, - *, - spec: DemoSpec, - scenario: Any, - adapter: DemoAdapter, - provider: Any, - ) -> None: - del spec, scenario, adapter, provider - - def create_run_context( - self, - *, - spec: DemoSpec, - adapter: DemoAdapter, - host: RuntimeHost, - model_warmup_plan: ModelWarmupPlan, - ) -> RunContext: - del spec, adapter - return RunContext( - host=host, - run_metrics=InMemorySessionMetricsRecorder(), - admission=SingleSessionAdmissionPolicy( - health_check=lambda: host.is_healthy, - ), - model_warmup_plan=model_warmup_plan, - ) - - def create_session_edges( - self, - *, - context: RunContext, - spec: DemoSpec, - scenario: Any, - provider: Any, - adapter: DemoAdapter, - ) -> SessionEdges: - del spec, scenario, provider, adapter - return SessionEdges( - input_source=_IOHandlerBatchInputSource(self.io_handler), - output_sink=_IOHandlerOutputSink(self.io_handler), - cleanup_tasks=context.cleanup_tasks, - transport=NoopTransportService(), - ) - - def select_driver(self) -> SessionDriver | AsyncSessionDriver: - return self.driver - - -@dataclass(slots=True) -class _IOHandlerBatchInputSource: - io_handler: IOHandler - is_finite: bool = False - is_deterministic: bool = False - user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) - - def is_finished(self) -> bool: - return self.io_handler.should_exit() - - def next_window(self, request: StepRequirements) -> UserInputWindow: - return self.io_handler.next_window(request) - - -@dataclass(slots=True) -class _IOHandlerOutputSink: - io_handler: IOHandler - produces_artifacts: bool = True - _generation_started: bool = field(default=False, init=False, repr=False) - - def open(self, session_info: SessionInfo) -> None: - self.io_handler.open(session_info) - - def begin_generation(self, generation: int) -> None: - self._generation_started = True - self.io_handler.begin_generation(generation) - - def write(self, result: StepResult) -> OutputDecision: - if not self._generation_started: - self.begin_generation(0) - return self.io_handler.emit_chunk(result) - - def close(self) -> Sequence[OutputArtifact]: - return self.io_handler.close() - - def _run_mode_is_async(run_mode: RunMode) -> bool: driver = run_mode.select_driver() return inspect.iscoroutinefunction(driver.run_one_session) diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 571b9f1d5..ab7e4c813 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -17,8 +17,14 @@ FrameOutputSink, InferenceSessionApplicationAdapter, IOHandler, + IOHandlerServer, + ReplayIOHandler, Runner, RuntimeOutputSinkFrameAdapter, + WebRTCIOHandlerServer, + create_native_window_io_handler, + create_replay_io_handler, + create_webrtc_io_handler, ) from flashdreams.runtime import ( CanonicalInputSchema, @@ -63,7 +69,13 @@ def test_public_demo_contracts_are_importable() -> None: assert Application.__name__ == "Application" assert ApplicationSession.__name__ == "ApplicationSession" assert IOHandler.__name__ == "IOHandler" + assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" + assert create_replay_io_handler.__name__ == "create_replay_io_handler" + assert create_native_window_io_handler.__name__ == ( + "create_native_window_io_handler" + ) + assert create_webrtc_io_handler.__name__ == "create_webrtc_io_handler" def test_inference_session_adapter_satisfies_application_session() -> None: @@ -226,6 +238,70 @@ async def test_runner_run_async_delegates_to_async_session_helper() -> None: assert io_handler.emitted_steps == [0] +def test_replay_io_factory_runs_through_public_runner() -> None: + output_sink = NullOutputSink(store_results=True) + io_handler = create_replay_io_handler(output_sink=output_sink) + app = _RunnerFakeApplication(total_steps=2) + + assert isinstance(io_handler, ReplayIOHandler) + result = Runner( + io_handler=io_handler, + app=app, + model_id="fake-replay-factory", + ).run() + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 2 + assert output_sink.output_count == 2 + + +def test_replay_io_factory_forwards_to_metric_tail() -> None: + metric_tail = _RecordingFrameOutputSink() + io_handler = create_replay_io_handler(metric_output_sink=metric_tail) + + result = Runner( + io_handler=io_handler, + app=_RunnerFakeApplication(total_steps=1), + model_id="fake-replay-factory", + ).run() + + assert result.status == "completed" + assert metric_tail.records == [(0.0, 0)] + + +def test_native_window_factory_exposes_io_handler_shape() -> None: + io_handler = create_native_window_io_handler((1280, 720)) + + assert isinstance(io_handler, IOHandler) + assert io_handler.run_mode.name == "native-window" + with pytest.raises(ValueError, match="viewport dimensions"): + create_native_window_io_handler((0, 720)) + + +def test_webrtc_io_factory_returns_server_not_ready_handler() -> None: + handler = _RecordingIOHandler() + server = create_webrtc_io_handler( + "127.0.0.1", + 8080, + (1280, 720), + handlers=(handler,), + ) + calls: list[IOHandler] = [] + + def run_session(io_handler: IOHandler) -> RunResult: + calls.append(io_handler) + return RunResult(status="completed") + + assert isinstance(server, IOHandlerServer) + assert isinstance(server, WebRTCIOHandlerServer) + result = server.serve(run_session) + + assert result.status == "completed" + assert calls == [handler] + assert not isinstance(server, IOHandler) + + class _FakeSession: inference_input_schema = InferenceInputSchema( step_fields=(InputField(name="chunk_index"),) @@ -344,6 +420,14 @@ def close(self) -> Sequence[OutputArtifact]: return () +class _RecordingFrameOutputSink: + def __init__(self) -> None: + self.records: list[tuple[float, int]] = [] + + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + self.records.append((timestamp_s, chunk.step_index)) + + class _RunnerFakeApplication: model_id = "runner-fake" From d118f9e55f4dd280c5b74f2f839d267c1a6ac24e Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 07:47:34 +0000 Subject: [PATCH 06/33] Adopt public demo Runner in app launch path Add the public DemoApplication app base and route replay launches through the public replay IO factory plus Runner. Keep the old runtime demo app import as a compatibility shim. Expose a server-shaped callback adapter for existing WebRTC serving code, and update OmniDreams and LingBot demo apps to use the new app-base IO selection hook instead of prepare_webrtc/serve_webrtc. Repoint non-benchmark finite flashdreams-run launches at the public replay Runner helper. Add CPU tests for app-base replay selection, server-shaped selection, and the adapter-backed Runner path. --- flashdreams/flashdreams/demo/__init__.py | 5 + flashdreams/flashdreams/demo/app.py | 119 +++++++++++++++ flashdreams/flashdreams/demo/application.py | 17 +++ flashdreams/flashdreams/demo/io.py | 20 +++ flashdreams/flashdreams/demo/runner.py | 79 ++++++++-- flashdreams/flashdreams/runtime/demo/app.py | 76 +--------- .../tests/test_demo_application_api.py | 137 +++++++++++++++++- flashdreams/tests/test_runtime_demo_api.py | 56 ++++++- integrations/lingbot/lingbot/demo/app.py | 44 +++--- integrations/lingbot/tests/test_demo_api.py | 16 +- .../omnidreams/omnidreams/demo/app.py | 26 ++-- .../omnidreams/tests/test_demo_api.py | 8 +- 12 files changed, 477 insertions(+), 126 deletions(-) create mode 100644 flashdreams/flashdreams/demo/app.py diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 798038f8f..7409dea0a 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -3,6 +3,7 @@ """Public demo application authoring API.""" +from flashdreams.demo.app import DemoApplication, run_replay_application from flashdreams.demo.application import ( Application, ApplicationSession, @@ -16,6 +17,7 @@ RuntimeOutputSinkFrameAdapter, ) from flashdreams.demo.io import ( + CallbackIOHandlerServer, IOHandlerServer, NativeWindowIOHandler, ReplayIOHandler, @@ -29,7 +31,9 @@ __all__ = [ "Application", "ApplicationSession", + "CallbackIOHandlerServer", "DemoAdapterApplication", + "DemoApplication", "FrameOutputSink", "IApplication", "IApplicationSession", @@ -45,4 +49,5 @@ "create_native_window_io_handler", "create_replay_io_handler", "create_webrtc_io_handler", + "run_replay_application", ] diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py new file mode 100644 index 000000000..4ba3bb134 --- /dev/null +++ b/flashdreams/flashdreams/demo/app.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared command lifecycle for public demo applications.""" + +from __future__ import annotations + +import argparse +import sys +from abc import ABC, abstractmethod +from typing import Any + +import torch +import torch.distributed as dist + +from flashdreams.core.distributed import init as distributed_init +from flashdreams.runtime.demo.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) +from flashdreams.runtime.demo.outputs import build_output_sink +from flashdreams.runtime.demo.run_modes import RunResult +from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec + +from .application import Application, DemoAdapterApplication, IOHandler +from .io import IOHandlerServer, create_replay_io_handler +from .runner import Runner + + +class DemoApplication(ABC): + """Base command application shared by model replay and WebRTC demos.""" + + def main(self, argv: list[str] | None = None) -> None: + """Parse arguments, select an IO bundle, and run the selected mode.""" + configure_logging() + args = self.parse_args(argv) + selection = self.create_io_handler(args) + + if isinstance(selection, IOHandlerServer): + result = selection.serve(lambda handler: self._run_handler(args, handler)) + else: + result = self._run_handler(args, selection) + _raise_for_failed_result(result) + + @abstractmethod + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + """Parse this model's command-line arguments.""" + + @abstractmethod + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + """Build the model-specific replay specification.""" + + @abstractmethod + def replay_adapter(self) -> DemoAdapter: + """Create the model-specific replay adapter.""" + + def application(self, args: argparse.Namespace) -> Application: + """Create the application object consumed by the public runner.""" + return DemoAdapterApplication( + adapter=self.replay_adapter(), + spec=self.replay_spec(args), + ) + + def create_io_handler( + self, + args: argparse.Namespace, + ) -> IOHandler | IOHandlerServer: + """Select the IO factory for the parsed command.""" + command = str(getattr(args, "command", "")) + if command == "replay": + spec = self.replay_spec(args) + return create_replay_io_handler(output_sink=build_output_sink(spec.output)) + if command == "webrtc": + context = initialize_cuda_distributed( + default_device=args.device, + distributed_init_fn=distributed_init, + configure_logging_fn=configure_logging, + torch_module=torch, + dist_module=dist, + ) + return self.webrtc_io_handler(args, context=context) + raise AssertionError(f"Unhandled command: {command}") + + def webrtc_io_handler( + self, + args: argparse.Namespace, + *, + context: Any, + ) -> IOHandlerServer: + """Create the WebRTC server-shaped IO factory for this model.""" + del args, context + raise ValueError("This demo application does not support WebRTC.") + + def _run_handler(self, args: argparse.Namespace, handler: IOHandler) -> RunResult: + return Runner( + io_handler=handler, + app=self.application(args), + ).run() + + +def run_replay_application(*, spec: DemoSpec, adapter: DemoAdapter) -> RunResult: + """Run a finite demo spec through the public replay IO factory and runner.""" + return Runner( + io_handler=create_replay_io_handler(output_sink=build_output_sink(spec.output)), + app=DemoAdapterApplication(adapter=adapter, spec=spec), + ).run() + + +def _raise_for_failed_result(result: RunResult) -> None: + if result.status in {"completed", "skipped"}: + return + reason = result.reason or (str(result.error) if result.error is not None else None) + if reason is None: + reason = f"Demo ended with status {result.status!r}." + print(reason, file=sys.stderr) + raise SystemExit(1) + + +__all__ = ["DemoApplication", "run_replay_application"] diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py index 8441e9d78..4e86957ff 100644 --- a/flashdreams/flashdreams/demo/application.py +++ b/flashdreams/flashdreams/demo/application.py @@ -223,6 +223,23 @@ def create_session(self) -> ApplicationSession: runtime.start_session(scenario.initial_inputs) ) + def start_session(self, inputs: InferenceInput) -> ApplicationSession: + """Create a runtime session from runner-provided initial model inputs.""" + scenario = self._scenario + if scenario is None: + self.init(()) + scenario = self._scenario + if scenario is None: + raise RuntimeError("DemoAdapterApplication failed to prepare a scenario.") + runtime = self.adapter.create_runtime(_require_config(self.spec)) + self._runtimes.append(runtime) + return InferenceSessionApplicationAdapter(runtime.start_session(inputs)) + + @property + def prepared_scenario(self) -> PreparedScenario | None: + """Return the scenario materialized by ``init(...)``, if any.""" + return self._scenario + def close(self) -> None: errors: list[Exception] = [] while self._runtimes: diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py index 181cffada..2498d66be 100644 --- a/flashdreams/flashdreams/demo/io.py +++ b/flashdreams/flashdreams/demo/io.py @@ -37,6 +37,7 @@ from .application import FrameOutputSink, IOHandler RunSessionCallback = Callable[[IOHandler], RunResult] +ServeCallback = Callable[[], object] @runtime_checkable @@ -176,6 +177,17 @@ def serve(self, run_session: RunSessionCallback) -> RunResult: return result +@dataclass(slots=True) +class CallbackIOHandlerServer: + """Server adapter for existing transports during IO-factory adoption.""" + + callback: ServeCallback + + def serve(self, run_session: RunSessionCallback) -> RunResult: + del run_session + return _coerce_run_result(self.callback()) + + @dataclass(slots=True) class IOHandlerRunMode: """Runtime run-mode adapter for public ``IOHandler`` instances.""" @@ -357,7 +369,14 @@ def _result_timestamp_s(result: StepResult) -> float: return result.output_window.start_s +def _coerce_run_result(value: object) -> RunResult: + if isinstance(value, RunResult): + return value + return RunResult(status="completed") + + __all__ = [ + "CallbackIOHandlerServer", "IOHandlerBatchInputSource", "IOHandlerOutputSink", "IOHandlerRunMode", @@ -365,6 +384,7 @@ def _result_timestamp_s(result: StepResult) -> float: "NativeWindowIOHandler", "ReplayIOHandler", "RunSessionCallback", + "ServeCallback", "WebRTCIOHandlerServer", "create_native_window_io_handler", "create_replay_io_handler", diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index cd898da33..85c687827 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -32,7 +32,12 @@ ProviderCapabilities, UserInputWindow, ) -from flashdreams.runtime.demo.spec import DemoSpec, NullOutputSpec, PreparedScenario +from flashdreams.runtime.demo.spec import ( + DemoAdapter, + DemoSpec, + NullOutputSpec, + PreparedScenario, +) from flashdreams.runtime.inputs import ( InferenceInput, InferenceInputSchema, @@ -41,7 +46,12 @@ from flashdreams.runtime.mapping import InputMapping from flashdreams.runtime.types import StepRequirements -from .application import Application, ApplicationSession, IOHandler +from .application import ( + Application, + ApplicationSession, + DemoAdapterApplication, + IOHandler, +) from .io import IOHandlerRunMode @@ -77,12 +87,12 @@ async def run_async(self) -> RunResult: return self._run_sync(run_mode=run_mode) host, owns_host = self._selected_host() - context = self._create_context(host) spec = self._create_spec(run_mode) - scenario = _runner_scenario() - adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + context = self._create_context(host) try: self.app.init(tuple(self.launch_args)) + scenario = self._create_scenario() + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) return await run_demo_session_async( context=context, spec=spec, @@ -99,12 +109,12 @@ async def run_async(self) -> RunResult: def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: selected_run_mode = run_mode or self._selected_run_mode() host, owns_host = self._selected_host() - context = self._create_context(host) spec = self._create_spec(selected_run_mode) - scenario = _runner_scenario() - adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + context = self._create_context(host) try: self.app.init(tuple(self.launch_args)) + scenario = self._create_scenario() + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) return run_demo_session( context=context, spec=spec, @@ -142,6 +152,8 @@ def _create_context(self, host: RuntimeHost) -> RunContext: ) def _create_spec(self, run_mode: RunMode) -> DemoSpec: + if isinstance(self.app, DemoAdapterApplication): + return self.app.spec model_id = self.model_id or _application_model_id(self.app) return DemoSpec( model_id=model_id, @@ -150,14 +162,27 @@ def _create_spec(self, run_mode: RunMode) -> DemoSpec: config=InferenceConfig(model_id=model_id), ) + def _create_scenario(self) -> PreparedScenario: + if isinstance(self.app, DemoAdapterApplication): + scenario = self.app.prepared_scenario + if scenario is None: + raise RuntimeError( + "DemoAdapterApplication did not prepare a scenario during init." + ) + return scenario + return _runner_scenario() + @dataclass(slots=True) class _ApplicationRuntime: app: Application def start_session(self, inputs: InferenceInput) -> InferenceSession: - del inputs - session = self.app.create_session() + start_session = getattr(self.app, "start_session", None) + if callable(start_session): + session = start_session(inputs) + else: + session = self.app.create_session() if not isinstance(session, ApplicationSession): raise TypeError( "Application.create_session() must return ApplicationSession, " @@ -182,16 +207,31 @@ def model_id(self) -> str: @property def inference_input_schema(self) -> InferenceInputSchema: + adapter = _application_adapter(self.app) + if adapter is not None: + return adapter.inference_input_schema return InferenceInputSchema() @property def canonical_input_schema(self) -> CanonicalInputSchema: + adapter = _application_adapter(self.app) + if adapter is not None and adapter.canonical_input_schema is not None: + return adapter.canonical_input_schema return CanonicalInputSchema() def default_input_mapping(self) -> InputMapping | None: + adapter = _application_adapter(self.app) + if adapter is not None: + default_input_mapping = getattr(adapter, "default_input_mapping", None) + if callable(default_input_mapping): + return default_input_mapping() return None def validate_config(self, config: InferenceConfig) -> None: + adapter = _application_adapter(self.app) + if adapter is not None: + adapter.validate_config(config) + return if config.model_id != self.spec.model_id: raise ValueError(f"Unsupported model_id={config.model_id!r}.") @@ -200,9 +240,15 @@ def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: return cast(InferenceRuntime, _ApplicationRuntime(self.app)) def supported_input_modes(self) -> tuple[str, ...]: + adapter = _application_adapter(self.app) + if adapter is not None: + return adapter.supported_input_modes() return (self.spec.input_mode,) def supported_output_modes(self) -> tuple[str, ...]: + adapter = _application_adapter(self.app) + if adapter is not None: + return adapter.supported_output_modes() return (self.spec.output.mode,) def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: @@ -214,7 +260,12 @@ def create_model_input_provider( self, spec: DemoSpec, scenario: Any, - ) -> "_RunnerModelInputProvider": + ) -> object: + adapter = _application_adapter(self.app) + if adapter is not None: + create_provider = getattr(adapter, "create_model_input_provider", None) + if callable(create_provider): + return create_provider(spec, scenario) del spec, scenario return _RunnerModelInputProvider() @@ -271,4 +322,10 @@ def _runner_scenario() -> PreparedScenario: return PreparedScenario(initial_inputs=InferenceInput()) +def _application_adapter(app: Application) -> DemoAdapter | None: + if isinstance(app, DemoAdapterApplication): + return app.adapter + return None + + __all__ = ["Runner"] diff --git a/flashdreams/flashdreams/runtime/demo/app.py b/flashdreams/flashdreams/runtime/demo/app.py index b763e4eaf..d91e6842e 100644 --- a/flashdreams/flashdreams/runtime/demo/app.py +++ b/flashdreams/flashdreams/runtime/demo/app.py @@ -1,80 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared command lifecycle for model demo applications.""" - -from __future__ import annotations - -import argparse -import sys -from abc import ABC, abstractmethod -from typing import Any - -import torch -import torch.distributed as dist - -from flashdreams.core.distributed import init as distributed_init -from flashdreams.runtime.demo.bootstrap import ( - configure_logging, - initialize_cuda_distributed, -) -from flashdreams.runtime.demo.replay import run_replay_demo -from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec - - -class DemoApplication(ABC): - """Base command application shared by model replay and WebRTC demos.""" - - def main(self, argv: list[str] | None = None) -> None: - """Parse arguments and dispatch the selected demo mode.""" - configure_logging() - args = self.parse_args(argv) - if args.command == "replay": - result = run_replay_demo( - spec=self.replay_spec(args), - adapter=self.replay_adapter(), - ) - if result.status != "completed": - reason = result.reason or ( - str(result.error) if result.error is not None else None - ) - if reason is None: - reason = f"Replay demo ended with status {result.status!r}." - print(reason, file=sys.stderr) - raise SystemExit(1) - return - if args.command == "webrtc": - context = initialize_cuda_distributed( - default_device=args.device, - distributed_init_fn=distributed_init, - configure_logging_fn=configure_logging, - torch_module=torch, - dist_module=dist, - ) - self.prepare_webrtc(args, context=context) - self.serve_webrtc(args, context=context) - return - raise AssertionError(f"Unhandled command: {args.command}") - - @abstractmethod - def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: - """Parse this model's command-line arguments.""" - - @abstractmethod - def replay_spec(self, args: argparse.Namespace) -> DemoSpec: - """Build the model-specific replay specification.""" - - @abstractmethod - def replay_adapter(self) -> DemoAdapter: - """Create the model-specific replay adapter.""" - - def prepare_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - """Perform optional model-specific setup before serving WebRTC.""" - del args, context - - @abstractmethod - def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - """Build and serve the model-specific WebRTC demo.""" +"""Compatibility import for the public demo application base.""" +from flashdreams.demo.app import DemoApplication __all__ = ["DemoApplication"] diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index ab7e4c813..6d6b8632a 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -3,6 +3,7 @@ from __future__ import annotations +import argparse import threading from collections.abc import Sequence from dataclasses import dataclass, field @@ -13,7 +14,9 @@ from flashdreams.demo import ( Application, ApplicationSession, + CallbackIOHandlerServer, DemoAdapterApplication, + DemoApplication, FrameOutputSink, InferenceSessionApplicationAdapter, IOHandler, @@ -50,6 +53,8 @@ NullOutputSpec, OutputDecision, PreparedScenario, + PreparedStep, + ProviderCapabilities, RunContext, RunModeCapabilities, RunResult, @@ -68,6 +73,7 @@ def test_public_demo_contracts_are_importable() -> None: assert Application.__name__ == "Application" assert ApplicationSession.__name__ == "ApplicationSession" + assert DemoApplication.__name__ == "DemoApplication" assert IOHandler.__name__ == "IOHandler" assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" @@ -302,6 +308,78 @@ def run_session(io_handler: IOHandler) -> RunResult: assert not isinstance(server, IOHandler) +def test_demo_application_replay_selects_factory_and_runner() -> None: + app = _ReplayDemoApplication() + + app.main(["replay"]) + + assert app.adapter.runtimes + runtime = app.adapter.runtimes[0] + assert runtime.session.closed + assert runtime.closed + + +def test_demo_application_server_selection_does_not_build_replay_app() -> None: + app = _ServerDemoApplication() + + app.main(["webrtc"]) + + assert app.server_called + assert not app.replay_adapter_called + + +class _ReplayDemoApplication(DemoApplication): + def __init__(self) -> None: + self.adapter = _FakeDemoAdapter() + + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("command", choices=("replay",)) + return parser.parse_args(argv) + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + assert args.command == "replay" + return DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ) + + def replay_adapter(self) -> "_FakeDemoAdapter": + return self.adapter + + +class _ServerDemoApplication(DemoApplication): + def __init__(self) -> None: + self.server_called = False + self.replay_adapter_called = False + + def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("command", choices=("webrtc",)) + return parser.parse_args(argv) + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + del args + raise AssertionError("server selection should not build a replay spec") + + def replay_adapter(self) -> "_FakeDemoAdapter": + self.replay_adapter_called = True + return _FakeDemoAdapter() + + def create_io_handler( + self, + args: argparse.Namespace, + ) -> IOHandler | IOHandlerServer: + assert args.command == "webrtc" + + def serve() -> RunResult: + self.server_called = True + return RunResult(status="completed") + + return CallbackIOHandlerServer(serve) + + class _FakeSession: inference_input_schema = InferenceInputSchema( step_fields=(InputField(name="chunk_index"),) @@ -311,6 +389,7 @@ def __init__(self) -> None: self.initialized = False self.closed = False self.reset_called = False + self.step_index = 0 def init(self) -> None: self.initialized = True @@ -319,19 +398,23 @@ def session_info(self) -> SessionInfo: return SessionInfo(output_layout="thwc") def next_step_request(self) -> StepRequest | None: + if self.step_index >= 1: + return None return StepRequest( - step_index=0, + step_index=self.step_index, inference_input_schema=self.inference_input_schema, ) def step(self, inputs: InferenceInput) -> StepResult: self.inference_input_schema.require_step(inputs) - return StepResult( - step_index=0, + result = StepResult( + step_index=self.step_index, output="chunk", frame_count=1, output_window=TimeWindow(start_s=0.0, end_s=1.0), ) + self.step_index += 1 + return result def reset(self, inputs: InferenceInput | None = None) -> None: del inputs @@ -392,6 +475,54 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: source_schema=UserInputSchema(), ) + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_FakeModelInputProvider": + assert spec.model_id == self.model_id + return _FakeModelInputProvider( + scenario=scenario, + inference_input_schema=self.inference_input_schema, + ) + + +class _FakeModelInputProvider: + def __init__( + self, + *, + scenario: PreparedScenario, + inference_input_schema: InferenceInputSchema, + ) -> None: + self._scenario = scenario + self.capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=inference_input_schema, + ) + self.closed = False + + def prepare_initial_input(self) -> InferenceInput: + return self._scenario.initial_inputs + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del user_window + return PreparedStep( + inference_input=InferenceInput(step={"chunk_index": request.step_index}) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.closed = True + class _FakeIOHandler: def open(self, session_info: SessionInfo) -> None: diff --git a/flashdreams/tests/test_runtime_demo_api.py b/flashdreams/tests/test_runtime_demo_api.py index 9600f19d0..74c5294eb 100644 --- a/flashdreams/tests/test_runtime_demo_api.py +++ b/flashdreams/tests/test_runtime_demo_api.py @@ -30,6 +30,7 @@ OutputArtifact, OutputTarget, StepRequest, + StepRequirements, StepResult, TimeWindow, UserInputs, @@ -44,7 +45,10 @@ OutputSink, OutputSpec, PreparedScenario, + PreparedStep, + ProviderCapabilities, RunResult, + UserInputWindow, WebRTCAppResources, WebRTCOutputSpec, build_output_sink, @@ -438,10 +442,6 @@ def replay_spec(self, args: argparse.Namespace) -> DemoSpec: def replay_adapter(self) -> "_FakeDemoAdapter": return self._adapter - def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - del args, context - raise AssertionError("webrtc should not run") - class _FakeDemoAdapter: model_id = "fake-demo" @@ -508,6 +508,54 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: raise ValueError("invalid scenario") return self.prepared_scenario + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_FakeModelInputProvider": + assert spec.model_id == self.model_id + return _FakeModelInputProvider( + scenario=scenario, + inference_input_schema=self.inference_input_schema, + ) + + +class _FakeModelInputProvider: + def __init__( + self, + *, + scenario: PreparedScenario, + inference_input_schema: InferenceInputSchema, + ) -> None: + self._scenario = scenario + self.capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + user_input_schema=scenario.source_schema, + inference_input_schema=inference_input_schema, + ) + self.closed = False + + def prepare_initial_input(self) -> InferenceInput: + return self._scenario.initial_inputs + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del user_window + return PreparedStep( + inference_input=InferenceInput(step={"chunk_index": request.step_index}) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.closed = True + class _FakeRuntime: def __init__( diff --git a/integrations/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py index c47abf9b9..d041a8a98 100644 --- a/integrations/lingbot/lingbot/demo/app.py +++ b/integrations/lingbot/lingbot/demo/app.py @@ -9,6 +9,8 @@ from pathlib import Path from typing import Any, Literal, cast +from flashdreams.demo import CallbackIOHandlerServer, IOHandlerServer +from flashdreams.demo.app import DemoApplication, run_replay_application from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( @@ -17,9 +19,7 @@ NullOutputSpec, WebRTCOutputSpec, ) -from flashdreams.runtime.demo.app import DemoApplication from flashdreams.runtime.demo.benchmark import run_benchmark_demo -from flashdreams.runtime.demo.replay import run_replay_demo from flashdreams.serving.webrtc.bootstrap import ( configure_logging, initialize_cuda_distributed, @@ -139,23 +139,29 @@ def replay_spec(self, args: argparse.Namespace) -> DemoSpec: def replay_adapter(self) -> LingbotDemoAdapter: return LingbotDemoAdapter() - def prepare_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - ensure_example_data_downloaded( - is_rank_zero=(context.world_rank == 0), - example_idx=args.example_idx, - ) - - def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - from .webrtc import serve_lingbot_webrtc_demo + def webrtc_io_handler( + self, + args: argparse.Namespace, + *, + context: Any, + ) -> IOHandlerServer: + def serve() -> object: + ensure_example_data_downloaded( + is_rank_zero=(context.world_rank == 0), + example_idx=args.example_idx, + ) + from .webrtc import serve_lingbot_webrtc_demo + + return serve_lingbot_webrtc_demo( + spec=_webrtc_spec( + args, + device=str(context.device), + context_parallel_size=context.world_size, + ), + world_rank=context.world_rank, + ) - serve_lingbot_webrtc_demo( - spec=_webrtc_spec( - args, - device=str(context.device), - context_parallel_size=context.world_size, - ), - world_rank=context.world_rank, - ) + return CallbackIOHandlerServer(serve) _APPLICATION = LingbotDemoApplication() @@ -230,7 +236,7 @@ def launch_from_runner( stats_dir=stats_dir, capture_output=True, ) - return run_replay_demo(spec=spec, adapter=LingbotDemoAdapter()) + return run_replay_application(spec=spec, adapter=LingbotDemoAdapter()) if mode != "webrtc": raise ValueError(f"Unsupported LingBot launch mode: {mode!r}.") diff --git a/integrations/lingbot/tests/test_demo_api.py b/integrations/lingbot/tests/test_demo_api.py index 929c36a10..cd7d0cda9 100644 --- a/integrations/lingbot/tests/test_demo_api.py +++ b/integrations/lingbot/tests/test_demo_api.py @@ -124,12 +124,16 @@ def test_lingbot_direct_runner_launch_builds_mp4_spec( ) -> None: captured: list[DemoSpec] = [] - def fake_run_replay_demo(*, spec: DemoSpec, adapter: object) -> str: + def fake_run_replay_application(*, spec: DemoSpec, adapter: object) -> str: del adapter captured.append(spec) return "completed" - monkeypatch.setattr(demo_app_module, "run_replay_demo", fake_run_replay_demo) + monkeypatch.setattr( + demo_app_module, + "run_replay_application", + fake_run_replay_application, + ) result = demo_app_module.launch_from_runner( config=RUNNER_LINGBOT_WORLD_FAST, @@ -188,12 +192,16 @@ def test_lingbot_direct_runner_launch_builds_null_spec( ) -> None: captured: list[DemoSpec] = [] - def fake_run_replay_demo(*, spec: DemoSpec, adapter: object) -> str: + def fake_run_replay_application(*, spec: DemoSpec, adapter: object) -> str: del adapter captured.append(spec) return "completed" - monkeypatch.setattr(demo_app_module, "run_replay_demo", fake_run_replay_demo) + monkeypatch.setattr( + demo_app_module, + "run_replay_application", + fake_run_replay_application, + ) result = demo_app_module.launch_from_runner( config=RUNNER_LINGBOT_WORLD_FAST, diff --git a/integrations/omnidreams/omnidreams/demo/app.py b/integrations/omnidreams/omnidreams/demo/app.py index 1ebb6038b..084e21996 100644 --- a/integrations/omnidreams/omnidreams/demo/app.py +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -13,6 +13,8 @@ from omnidreams.runner import DEFAULT_EXAMPLE_DATA_UUID_1V +from flashdreams.demo import CallbackIOHandlerServer, IOHandlerServer +from flashdreams.demo.app import DemoApplication, run_replay_application from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( @@ -21,9 +23,7 @@ NullOutputSpec, WebRTCOutputSpec, ) -from flashdreams.runtime.demo.app import DemoApplication from flashdreams.runtime.demo.benchmark import run_benchmark_demo -from flashdreams.runtime.demo.replay import run_replay_demo from flashdreams.serving.webrtc.bootstrap import ( configure_logging, initialize_cuda_distributed, @@ -137,13 +137,21 @@ def replay_spec(self, args: argparse.Namespace) -> DemoSpec: def replay_adapter(self) -> OmnidreamsDemoAdapter: return OmnidreamsDemoAdapter() - def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: - from .webrtc import serve_omnidreams_webrtc_demo + def webrtc_io_handler( + self, + args: argparse.Namespace, + *, + context: Any, + ) -> IOHandlerServer: + def serve() -> object: + from .webrtc import serve_omnidreams_webrtc_demo + + return serve_omnidreams_webrtc_demo( + spec=_webrtc_spec(args, device=str(context.device)), + world_rank=context.world_rank, + ) - serve_omnidreams_webrtc_demo( - spec=_webrtc_spec(args, device=str(context.device)), - world_rank=context.world_rank, - ) + return CallbackIOHandlerServer(serve) _APPLICATION = OmnidreamsDemoApplication() @@ -226,7 +234,7 @@ def launch_from_runner( stats_dir=stats_dir, capture_output=True, ) - return run_replay_demo(spec=spec, adapter=OmnidreamsDemoAdapter()) + return run_replay_application(spec=spec, adapter=OmnidreamsDemoAdapter()) if mode != "webrtc": raise ValueError(f"Unsupported OmniDreams launch mode: {mode!r}.") diff --git a/integrations/omnidreams/tests/test_demo_api.py b/integrations/omnidreams/tests/test_demo_api.py index 09b51994f..ee4112512 100644 --- a/integrations/omnidreams/tests/test_demo_api.py +++ b/integrations/omnidreams/tests/test_demo_api.py @@ -101,12 +101,16 @@ def test_omnidreams_direct_runner_launch_builds_null_spec( ) -> None: captured: list[DemoSpec] = [] - def fake_run_replay_demo(*, spec: DemoSpec, adapter: object) -> str: + def fake_run_replay_application(*, spec: DemoSpec, adapter: object) -> str: del adapter captured.append(spec) return "completed" - monkeypatch.setattr(demo_app_module, "run_replay_demo", fake_run_replay_demo) + monkeypatch.setattr( + demo_app_module, + "run_replay_application", + fake_run_replay_application, + ) config = OMNIDREAMS_RUNNERS["omnidreams"] result = demo_app_module.launch_from_runner( From 433fd037494232cdd3cf284e8120f39eb7e3b477 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:01:14 +0000 Subject: [PATCH 07/33] Ensure public demo Runner always closes applications Make Runner own Application.close() explicitly instead of relying on owned host cleanup to reach it. Application cleanup now runs through the host worker for both owned and externally supplied hosts, and is attempted even when context cleanup raises. Add focused sync and async coverage for external-host cleanup and cleanup failure paths. --- flashdreams/flashdreams/demo/runner.py | 77 +++++++++++++++++-- .../tests/test_demo_application_api.py | 69 +++++++++++++++++ 2 files changed, 139 insertions(+), 7 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 85c687827..5bd814de8 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -102,9 +102,12 @@ async def run_async(self) -> RunResult: pipeline=self.pipeline or StepPipeline(), ) finally: - await context.close_async() - if owns_host: - host.close() + await _close_runner_resources_async( + context=context, + host=host, + app=self.app, + owns_host=owns_host, + ) def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: selected_run_mode = run_mode or self._selected_run_mode() @@ -124,9 +127,12 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: pipeline=self.pipeline or StepPipeline(), ) finally: - context.close() - if owns_host: - host.close() + _close_runner_resources( + context=context, + host=host, + app=self.app, + owns_host=owns_host, + ) def _selected_run_mode(self) -> RunMode: if self.run_mode is not None: @@ -192,7 +198,7 @@ def start_session(self, inputs: InferenceInput) -> InferenceSession: return cast(InferenceSession, session) def close(self) -> None: - self.app.close() + return None @dataclass(slots=True) @@ -322,6 +328,63 @@ def _runner_scenario() -> PreparedScenario: return PreparedScenario(initial_inputs=InferenceInput()) +def _close_runner_resources( + *, + context: RunContext, + host: RuntimeHost, + app: Application, + owns_host: bool, +) -> None: + errors: list[Exception] = [] + _record_cleanup_error(errors, context.close) + _record_cleanup_error(errors, host.call, app.close) + if owns_host: + _record_cleanup_error(errors, host.close) + _raise_first_cleanup_error(errors) + + +async def _close_runner_resources_async( + *, + context: RunContext, + host: RuntimeHost, + app: Application, + owns_host: bool, +) -> None: + errors: list[Exception] = [] + try: + await context.close_async() + except Exception as exc: + errors.append(exc) + try: + await host.call_async(app.close) + except Exception as exc: + errors.append(exc) + if owns_host: + _record_cleanup_error(errors, host.close) + _raise_first_cleanup_error(errors) + + +def _record_cleanup_error( + errors: list[Exception], + cleanup: Any, + /, + *args: Any, +) -> None: + try: + cleanup(*args) + except Exception as exc: + errors.append(exc) + + +def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: + if not errors: + return + first = errors[0] + for extra in errors[1:]: + first.add_note(f"Additional cleanup error: {extra!r}") + raise first + + def _application_adapter(app: Application) -> DemoAdapter | None: if isinstance(app, DemoAdapterApplication): return app.adapter diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 6d6b8632a..533466e13 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -58,6 +58,7 @@ RunContext, RunModeCapabilities, RunResult, + RuntimeHost, SessionEdges, SessionInfo, SingleSessionAdmissionPolicy, @@ -224,6 +225,38 @@ def test_runner_run_drives_public_app_through_shared_runtime_path() -> None: assert io_handler.closed +def test_runner_closes_public_app_when_host_is_external() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = RuntimeHost(_ExternalApplicationRuntime(app)) + + try: + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + model_id="fake-runner", + ).run() + finally: + host.close() + + assert result.status == "completed" + assert app.closed + + +def test_runner_closes_public_app_when_context_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1) + + with pytest.raises(RuntimeError, match="run metrics close failed"): + Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + model_id="fake-runner", + ).run() + + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_delegates_to_async_session_helper() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -244,6 +277,22 @@ async def test_runner_run_async_delegates_to_async_session_helper() -> None: assert io_handler.emitted_steps == [0] +@pytest.mark.asyncio +async def test_runner_run_async_closes_public_app_when_context_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1) + + with pytest.raises(RuntimeError, match="run metrics close failed"): + await Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + + assert app.closed + + def test_replay_io_factory_runs_through_public_runner() -> None: output_sink = NullOutputSink(store_results=True) io_handler = create_replay_io_handler(output_sink=output_sink) @@ -581,6 +630,26 @@ def close(self) -> None: self.closed = True +class _ExternalApplicationRuntime: + def __init__(self, app: _RunnerFakeApplication) -> None: + self._app = app + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + session = self._app.create_session() + session.init() + return session + + def close(self) -> None: + self.closed = True + + +class _FailingCloseMetricsRecorder(InMemorySessionMetricsRecorder): + def close(self) -> Any: + raise RuntimeError("run metrics close failed") + + class _RunnerFakeSession: def __init__(self, *, total_steps: int) -> None: self.total_steps = total_steps From 3b56a0f01d5620aac0405a6674f497d6f66a5a58 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:06:07 +0000 Subject: [PATCH 08/33] Fix public demo cleanup type-check issues Make cleanup note handling type-checker friendly and update the fake public demo session to satisfy the runtime session protocol through a test adapter. --- flashdreams/flashdreams/demo/runner.py | 4 ++- .../tests/test_demo_application_api.py | 25 ++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 5bd814de8..9c7aae560 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -380,8 +380,10 @@ def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: if not errors: return first = errors[0] + add_note = getattr(first, "add_note", None) for extra in errors[1:]: - first.add_note(f"Additional cleanup error: {extra!r}") + if callable(add_note): + add_note(f"Additional cleanup error: {extra!r}") raise first diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 533466e13..4b54df79a 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -639,7 +639,7 @@ def start_session(self, inputs: InferenceInput) -> InferenceSession: del inputs session = self._app.create_session() session.init() - return session + return _ExternalInferenceSession(session) def close(self) -> None: self.closed = True @@ -650,6 +650,29 @@ def close(self) -> Any: raise RuntimeError("run metrics close failed") +class _ExternalInferenceSession: + def __init__(self, session: "_RunnerFakeSession") -> None: + self._session = session + + def session_info(self) -> SessionInfo: + return self._session.session_info() + + def next_step_request(self) -> StepRequest | None: + requirements = self._session.next_step_requirements() + if requirements is None: + return None + return StepRequest(step_index=requirements.step_index) + + def step(self, inputs: InferenceInput) -> StepResult: + return self._session.step(inputs) + + def reset(self, inputs: InferenceInput | None = None) -> None: + self._session.reset(inputs) + + def close(self) -> None: + self._session.close() + + class _RunnerFakeSession: def __init__(self, *, total_steps: int) -> None: self.total_steps = total_steps From 9bea1500083dfda357ca645d7214e377f6a44c98 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:15:53 +0000 Subject: [PATCH 09/33] Preserve runner failures during cleanup --- flashdreams/flashdreams/demo/runner.py | 51 +++++++++-- .../tests/test_demo_application_api.py | 90 ++++++++++++++++++- 2 files changed, 133 insertions(+), 8 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 9c7aae560..61167b643 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -89,11 +89,13 @@ async def run_async(self) -> RunResult: host, owns_host = self._selected_host() spec = self._create_spec(run_mode) context = self._create_context(host) + result: RunResult | None = None + primary_error: Exception | None = None try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) - return await run_demo_session_async( + result = await run_demo_session_async( context=context, spec=spec, scenario=scenario, @@ -101,12 +103,18 @@ async def run_async(self) -> RunResult: run_mode=run_mode, pipeline=self.pipeline or StepPipeline(), ) + return result + except Exception as exc: + primary_error = exc + raise finally: await _close_runner_resources_async( context=context, host=host, app=self.app, owns_host=owns_host, + run_result=result, + primary_error=primary_error, ) def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: @@ -114,11 +122,13 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: host, owns_host = self._selected_host() spec = self._create_spec(selected_run_mode) context = self._create_context(host) + result: RunResult | None = None + primary_error: Exception | None = None try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) - return run_demo_session( + result = run_demo_session( context=context, spec=spec, scenario=scenario, @@ -126,12 +136,18 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: run_mode=selected_run_mode, pipeline=self.pipeline or StepPipeline(), ) + return result + except Exception as exc: + primary_error = exc + raise finally: _close_runner_resources( context=context, host=host, app=self.app, owns_host=owns_host, + run_result=result, + primary_error=primary_error, ) def _selected_run_mode(self) -> RunMode: @@ -334,12 +350,20 @@ def _close_runner_resources( host: RuntimeHost, app: Application, owns_host: bool, + run_result: RunResult | None, + primary_error: Exception | None, ) -> None: errors: list[Exception] = [] _record_cleanup_error(errors, context.close) _record_cleanup_error(errors, host.call, app.close) if owns_host: _record_cleanup_error(errors, host.close) + if primary_error is not None: + _record_cleanup_notes(primary_error, errors) + return + if run_result is not None and run_result.status == "failed": + _record_cleanup_notes(run_result.error, errors) + return _raise_first_cleanup_error(errors) @@ -349,6 +373,8 @@ async def _close_runner_resources_async( host: RuntimeHost, app: Application, owns_host: bool, + run_result: RunResult | None, + primary_error: Exception | None, ) -> None: errors: list[Exception] = [] try: @@ -361,6 +387,12 @@ async def _close_runner_resources_async( errors.append(exc) if owns_host: _record_cleanup_error(errors, host.close) + if primary_error is not None: + _record_cleanup_notes(primary_error, errors) + return + if run_result is not None and run_result.status == "failed": + _record_cleanup_notes(run_result.error, errors) + return _raise_first_cleanup_error(errors) @@ -380,11 +412,20 @@ def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: if not errors: return first = errors[0] - add_note = getattr(first, "add_note", None) - for extra in errors[1:]: + _record_cleanup_notes(first, errors[1:]) + raise first + + +def _record_cleanup_notes( + primary: Exception | None, + errors: Sequence[Exception], +) -> None: + if primary is None: + return + add_note = getattr(primary, "add_note", None) + for extra in errors: if callable(add_note): add_note(f"Additional cleanup error: {extra!r}") - raise first def _application_adapter(app: Application) -> DemoAdapter | None: diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 4b54df79a..4d315843a 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -257,6 +257,38 @@ def test_runner_closes_public_app_when_context_cleanup_fails() -> None: assert app.closed +def test_runner_preserves_primary_error_when_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1, fail_init=True) + + with pytest.raises(RuntimeError, match="fake init failed") as exc_info: + Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + model_id="fake-runner", + ).run() + + notes = getattr(exc_info.value, "__notes__", ()) + assert any("run metrics close failed" in note for note in notes) + assert app.closed + + +def test_runner_preserves_failed_result_when_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1, fail_step=0) + + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + model_id="fake-runner", + ).run() + + assert result.status == "failed" + assert result.error is not None + assert str(result.error) == "fake step failed" + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_delegates_to_async_session_helper() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -293,6 +325,42 @@ async def test_runner_run_async_closes_public_app_when_context_cleanup_fails() - assert app.closed +@pytest.mark.asyncio +async def test_runner_run_async_preserves_primary_error_when_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1, fail_init=True) + + with pytest.raises(RuntimeError, match="fake init failed") as exc_info: + await Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + + notes = getattr(exc_info.value, "__notes__", ()) + assert any("run metrics close failed" in note for note in notes) + assert app.closed + + +@pytest.mark.asyncio +async def test_runner_run_async_preserves_failed_result_when_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1, fail_step=0) + + result = await Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + + assert result.status == "failed" + assert result.error is not None + assert str(result.error) == "fake step failed" + assert app.closed + + def test_replay_io_factory_runs_through_public_runner() -> None: output_sink = NullOutputSink(store_results=True) io_handler = create_replay_io_handler(output_sink=output_sink) @@ -611,8 +679,16 @@ def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: class _RunnerFakeApplication: model_id = "runner-fake" - def __init__(self, *, total_steps: int) -> None: + def __init__( + self, + *, + total_steps: int, + fail_init: bool = False, + fail_step: int | None = None, + ) -> None: self.total_steps = total_steps + self.fail_init = fail_init + self.fail_step = fail_step self.launch_args: tuple[str, ...] = () self.init_thread_id: int | None = None self.session: _RunnerFakeSession | None = None @@ -621,9 +697,14 @@ def __init__(self, *, total_steps: int) -> None: def init(self, launch_args: Sequence[str]) -> None: self.launch_args = tuple(launch_args) self.init_thread_id = threading.get_ident() + if self.fail_init: + raise RuntimeError("fake init failed") def create_session(self) -> "_RunnerFakeSession": - self.session = _RunnerFakeSession(total_steps=self.total_steps) + self.session = _RunnerFakeSession( + total_steps=self.total_steps, + fail_step=self.fail_step, + ) return self.session def close(self) -> None: @@ -674,8 +755,9 @@ def close(self) -> None: class _RunnerFakeSession: - def __init__(self, *, total_steps: int) -> None: + def __init__(self, *, total_steps: int, fail_step: int | None) -> None: self.total_steps = total_steps + self.fail_step = fail_step self.step_index = 0 self.init_thread_id: int | None = None self.step_thread_ids: tuple[int, ...] = () @@ -695,6 +777,8 @@ def next_step_requirements(self) -> StepRequirements | None: def step(self, model_input: InferenceInput) -> StepResult: assert model_input.step["step_index"] == self.step_index assert isinstance(model_input.step["user_window"], UserInputWindow) + if self.fail_step == self.step_index: + raise RuntimeError("fake step failed") self.step_thread_ids = (*self.step_thread_ids, threading.get_ident()) result = StepResult( step_index=self.step_index, From e4775468cfd1e9908536cd31746ee53dfff9cbcf Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:25:47 +0000 Subject: [PATCH 10/33] Unify demo session stop conditions --- flashdreams/flashdreams/demo/io.py | 6 +- .../flashdreams/runtime/demo/__init__.py | 4 + .../flashdreams/runtime/demo/drivers.py | 35 ++++---- .../flashdreams/runtime/demo/run_modes.py | 82 ++++++++++++++++++- .../tests/test_demo_application_api.py | 37 +++++++++ .../test_demo_runtime_realtime_driver.py | 4 + .../tests/test_demo_runtime_vertical_slice.py | 53 ++++++++++++ 7 files changed, 203 insertions(+), 18 deletions(-) diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py index 2498d66be..79fd650f9 100644 --- a/flashdreams/flashdreams/demo/io.py +++ b/flashdreams/flashdreams/demo/io.py @@ -67,6 +67,7 @@ class ReplayIOHandler: repr=False, ) _generation: int | None = field(default=None, init=False, repr=False) + _should_exit: bool = field(default=False, init=False, repr=False) _closed: bool = field(default=False, init=False, repr=False) def __post_init__(self) -> None: @@ -120,13 +121,16 @@ def emit_chunk(self, result: StepResult) -> OutputDecision: decision = OutputDecision() if self.metric_output_sink is not None: self.metric_output_sink.handle_output(_result_timestamp_s(result), result) + if decision.should_stop: + self._should_exit = True return decision def should_exit(self) -> bool: - return self._closed + return self._should_exit or self._closed def close(self) -> Sequence[OutputArtifact]: self._closed = True + self._should_exit = True close = getattr(self._output_sink, "close", None) if callable(close): artifacts = close() diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index bf64905b2..b3d6598dd 100644 --- a/flashdreams/flashdreams/runtime/demo/__init__.py +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -57,6 +57,8 @@ RunSummary, SessionDriver, SessionEdges, + SessionExitSource, + SessionExitState, SingleSessionAdmissionPolicy, WebRTCErrorPolicy, build_model_warmup_plan, @@ -175,6 +177,8 @@ "RuntimeHost", "SessionEdges", "SessionDriver", + "SessionExitState", + "SessionExitSource", "SessionInfo", "SignalActivationPolicy", "SingleSessionAdmissionPolicy", diff --git a/flashdreams/flashdreams/runtime/demo/drivers.py b/flashdreams/flashdreams/runtime/demo/drivers.py index 27179ac32..000c244f2 100644 --- a/flashdreams/flashdreams/runtime/demo/drivers.py +++ b/flashdreams/flashdreams/runtime/demo/drivers.py @@ -83,7 +83,7 @@ def run_one_session( if session is None: raise DriverInvariantError("setup_ok was set without a session.") try: - if session_edges.input_source.is_finished(): + if session_edges.should_exit(): break request = _next_step_requirements(host=host, session=session) if request is None: @@ -103,9 +103,9 @@ def run_one_session( if not outcome.control.provider_already_reset: host.call(provider.reset, outcome.control.reset_input) continue - if outcome.control.close_session: - break - if outcome.output.should_stop: + session_edges.observe_control(outcome.control) + session_edges.observe_output(outcome.output) + if session_edges.should_exit(poll_edges=False): break except DriverInvariantError: raise @@ -195,9 +195,9 @@ async def run_one_session( if not activation_result.activated: final_status = "not_activated" final_reason = activation_result.reason - elif not session_edges.transport.is_active(): + elif session_edges.should_exit(): final_status = "not_activated" - final_reason = "transport closed before first step" + final_reason = _not_activated_exit_reason(session_edges) else: try: initial_input = await host.call_async( @@ -224,10 +224,10 @@ async def run_one_session( while setup_ok: if session is None: raise DriverInvariantError("setup_ok was set without a session.") - if not session_edges.transport.is_active(): + if session_edges.should_exit(): if not first_step_started: final_status = "not_activated" - final_reason = "transport closed before first step" + final_reason = _not_activated_exit_reason(session_edges) break try: request = await _next_step_requirements_async( @@ -241,12 +241,9 @@ async def run_one_session( clock=clock, ) session_edges.metrics.record_catch_up(window_result.catch_up) - if ( - not session_edges.transport.is_active() - and not first_step_started - ): + if session_edges.should_exit() and not first_step_started: final_status = "not_activated" - final_reason = "transport closed before first step" + final_reason = _not_activated_exit_reason(session_edges) break outcome = await host.call_async( pipeline.execute_step, @@ -271,9 +268,9 @@ async def run_one_session( generation += 1 session_edges.output_sink.begin_generation(generation) continue - if outcome.control.close_session: - break - if outcome.output.should_stop: + session_edges.observe_control(outcome.control) + session_edges.observe_output(outcome.output) + if session_edges.should_exit(poll_edges=False): break if outcome.output.backpressure_s > 0: await clock.apply_backpressure(outcome.output.backpressure_s) @@ -329,6 +326,12 @@ def _mark_host_cleanup_failed(host: RuntimeHost, exc: Exception | None = None) - host.mark_unhealthy(_MODEL_CLEANUP_FAILED_REASON, exc) +def _not_activated_exit_reason(session_edges: SessionEdges) -> str: + if session_edges.exit_state.source == "transport_closed": + return "transport closed before first step" + return session_edges.exit_state.reason or "session stopped before first step" + + def run_demo_session( *, context: RunContext, diff --git a/flashdreams/flashdreams/runtime/demo/run_modes.py b/flashdreams/flashdreams/runtime/demo/run_modes.py index cb93e4e7f..01b589990 100644 --- a/flashdreams/flashdreams/runtime/demo/run_modes.py +++ b/flashdreams/flashdreams/runtime/demo/run_modes.py @@ -20,7 +20,8 @@ from flashdreams.runtime.output import OutputArtifact from .host import ModelWarmupPlan, WarmupSessionInputs -from .outputs import OutputSink +from .outputs import OutputDecision, OutputSink +from .session_inputs import ControlDecision if TYPE_CHECKING: from .host import RuntimeHost @@ -46,6 +47,22 @@ "not_activated", ] +SessionExitSource = Literal[ + "control_close", + "input_finished", + "output_stop", + "transport_closed", +] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class SessionExitState: + """Unified normal-stop state for a session loop.""" + + should_exit: bool = False + source: SessionExitSource | None = None + reason: str | None = None + @dataclass(frozen=True, kw_only=True, slots=True) class RunResult: @@ -328,12 +345,73 @@ class SessionEdges: clock: "RealtimeClock | DeterministicClock | None" = None activation: "ActivationPolicy | None" = None _closed_result: RunResult | None = field(default=None, init=False, repr=False) + _exit_state: SessionExitState = field( + default_factory=SessionExitState, + init=False, + repr=False, + ) @property def is_closed(self) -> bool: """Return whether ``close_result(...)`` has already finalized this session.""" return self._closed_result is not None + @property + def exit_state(self) -> SessionExitState: + """Return the currently observed normal-stop state.""" + return self._exit_state + + def should_exit(self, *, poll_edges: bool = True) -> bool: + """Return whether the driver should stop requesting model steps.""" + if self._exit_state.should_exit: + return True + if self.is_closed: + return True + if not poll_edges: + return False + if self.input_source.is_finished(): + self.request_exit( + source="input_finished", + reason="input source finished", + ) + return True + if not self.transport.is_active(): + reason = getattr(self.transport, "close_reason", None) + self.request_exit( + source="transport_closed", + reason=reason if isinstance(reason, str) else "transport closed", + ) + return True + return False + + def request_exit( + self, + *, + source: SessionExitSource, + reason: str | None = None, + ) -> None: + """Record a normal-stop request without replacing the first source.""" + if self._exit_state.should_exit: + return + self._exit_state = SessionExitState( + should_exit=True, + source=source, + reason=reason, + ) + + def observe_control(self, control: ControlDecision) -> None: + """Record provider-authored close decisions in the shared stop state.""" + if control.close_session: + self.request_exit( + source="control_close", + reason=control.reason or "control requested session close", + ) + + def observe_output(self, decision: OutputDecision) -> None: + """Record output-authored stop decisions in the shared stop state.""" + if decision.should_stop: + self.request_exit(source="output_stop", reason="output requested stop") + def record_cleanup_error(self, exc: Exception) -> None: """Record a cleanup error without letting metrics failures block teardown.""" try: @@ -524,6 +602,8 @@ def _coerce_warmup_sessions(value: object) -> tuple[WarmupSessionInputs, ...]: "RunSummary", "SessionEdges", "SessionDriver", + "SessionExitSource", + "SessionExitState", "SessionMetricsRecorder", "SessionReservation", "SessionStatus", diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 4d315843a..27f4c9f39 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -379,6 +379,23 @@ def test_replay_io_factory_runs_through_public_runner() -> None: assert output_sink.output_count == 2 +def test_replay_io_factory_should_exit_tracks_output_stop() -> None: + output_sink = _StoppingOutputSink() + io_handler = create_replay_io_handler(output_sink=output_sink) + + result = Runner( + io_handler=io_handler, + app=_RunnerFakeApplication(total_steps=5), + model_id="fake-replay-factory", + ).run() + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert output_sink.results == [0] + assert io_handler.should_exit() + + def test_replay_io_factory_forwards_to_metric_tail() -> None: metric_tail = _RecordingFrameOutputSink() io_handler = create_replay_io_handler(metric_output_sink=metric_tail) @@ -676,6 +693,26 @@ def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: self.records.append((timestamp_s, chunk.step_index)) +class _StoppingOutputSink: + produces_artifacts = False + + def __init__(self) -> None: + self.results: list[int] = [] + + def open(self, session_info: SessionInfo) -> None: + del session_info + + def begin_generation(self, generation: int) -> None: + del generation + + def write(self, result: StepResult) -> OutputDecision: + self.results.append(result.step_index) + return OutputDecision(should_stop=True) + + def close(self) -> Sequence[OutputArtifact]: + return () + + class _RunnerFakeApplication: model_id = "runner-fake" diff --git a/flashdreams/tests/test_demo_runtime_realtime_driver.py b/flashdreams/tests/test_demo_runtime_realtime_driver.py index 7195218ec..a98671357 100644 --- a/flashdreams/tests/test_demo_runtime_realtime_driver.py +++ b/flashdreams/tests/test_demo_runtime_realtime_driver.py @@ -104,6 +104,8 @@ async def test_realtime_driver_transport_close_before_first_step_is_not_activate assert result.status == "not_activated" assert result.reason == "transport closed before first step" assert session.step_inputs == [] + assert edges.exit_state.source == "transport_closed" + assert edges.should_exit() @pytest.mark.asyncio @@ -441,6 +443,8 @@ async def test_realtime_driver_applies_backpressure_through_clock() -> None: assert clock.backpressure == [0.25] assert metrics.catch_up_count == 2 assert len(output.results) == 2 + assert edges.exit_state.source == "output_stop" + assert edges.should_exit() @pytest.mark.asyncio diff --git a/flashdreams/tests/test_demo_runtime_vertical_slice.py b/flashdreams/tests/test_demo_runtime_vertical_slice.py index 826684033..468b6e63a 100644 --- a/flashdreams/tests/test_demo_runtime_vertical_slice.py +++ b/flashdreams/tests/test_demo_runtime_vertical_slice.py @@ -191,6 +191,59 @@ def test_batch_driver_slices_windows_from_step_requirements() -> None: ] +def test_batch_driver_stops_on_input_finished_through_session_edges() -> None: + session = _FakeVideoSession(num_steps=5) + runtime = _FakeVideoRuntime(session=session) + host = _RecordingRuntimeHost(runtime) + input_source = _FakeBatchInputSource(num_windows=1) + output = _RecordingOutputSink() + edges = SessionEdges( + input_source=input_source, + output_sink=output, + cleanup_tasks=set(), + metrics=InMemorySessionMetricsRecorder(), + ) + + result = BatchSessionDriver().run_one_session( + host=host, + provider=_FakeVideoModelInputProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert edges.exit_state.source == "input_finished" + assert edges.should_exit() + + +def test_batch_driver_stops_on_output_decision_through_session_edges() -> None: + session = _FakeVideoSession(num_steps=5) + runtime = _FakeVideoRuntime(session=session) + host = _RecordingRuntimeHost(runtime) + output = _RecordingOutputSink(decision=OutputDecision(should_stop=True)) + edges = SessionEdges( + input_source=_FakeBatchInputSource(num_windows=5), + output_sink=output, + cleanup_tasks=set(), + metrics=InMemorySessionMetricsRecorder(), + ) + + result = BatchSessionDriver().run_one_session( + host=host, + provider=_FakeVideoModelInputProvider(), + session_edges=edges, + pipeline=StepPipeline(), + ) + + assert result.status == "completed" + assert result.metrics is not None + assert result.metrics.counters["steps"] == 1 + assert edges.exit_state.source == "output_stop" + assert edges.should_exit() + + def test_run_demo_session_builds_edges_and_records_session_once() -> None: session = _FakeVideoSession(num_steps=1) runtime = _FakeVideoRuntime(session=session) From baf857895b69e669ebdfc4d70700c4c33811f077 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:34:54 +0000 Subject: [PATCH 11/33] Harden public demo runner cleanup --- flashdreams/flashdreams/demo/runner.py | 51 ++++++- flashdreams/flashdreams/runtime/demo/host.py | 5 + .../tests/test_demo_application_api.py | 125 +++++++++++++++++- 3 files changed, 174 insertions(+), 7 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 61167b643..0aac516f2 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -355,7 +355,7 @@ def _close_runner_resources( ) -> None: errors: list[Exception] = [] _record_cleanup_error(errors, context.close) - _record_cleanup_error(errors, host.call, app.close) + _close_application(errors=errors, host=host, app=app) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: @@ -381,10 +381,7 @@ async def _close_runner_resources_async( await context.close_async() except Exception as exc: errors.append(exc) - try: - await host.call_async(app.close) - except Exception as exc: - errors.append(exc) + await _close_application_async(errors=errors, host=host, app=app) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: @@ -408,6 +405,50 @@ def _record_cleanup_error( errors.append(exc) +def _close_application( + *, + errors: list[Exception], + host: RuntimeHost, + app: Application, +) -> None: + invoked = False + + def close_app() -> None: + nonlocal invoked + invoked = True + app.close() + + try: + host.call(close_app) + except Exception as exc: + if not invoked and host.is_closed: + _record_cleanup_error(errors, app.close) + return + errors.append(exc) + + +async def _close_application_async( + *, + errors: list[Exception], + host: RuntimeHost, + app: Application, +) -> None: + invoked = False + + def close_app() -> None: + nonlocal invoked + invoked = True + app.close() + + try: + await host.call_async(close_app) + except Exception as exc: + if not invoked and host.is_closed: + _record_cleanup_error(errors, app.close) + return + errors.append(exc) + + def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: if not errors: return diff --git a/flashdreams/flashdreams/runtime/demo/host.py b/flashdreams/flashdreams/runtime/demo/host.py index 6a6e93ce5..787e07147 100644 --- a/flashdreams/flashdreams/runtime/demo/host.py +++ b/flashdreams/flashdreams/runtime/demo/host.py @@ -81,6 +81,11 @@ def is_healthy(self) -> bool: """Return whether admission should continue accepting sessions.""" return self._healthy and not self._closed + @property + def is_closed(self) -> bool: + """Return whether the host has stopped accepting dispatched work.""" + return self._closed + @property def unhealthy_reason(self) -> str | None: """Return the first latched unhealthy reason, if any.""" diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 27f4c9f39..a26ccc32d 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -243,6 +243,24 @@ def test_runner_closes_public_app_when_host_is_external() -> None: assert app.closed +def test_runner_closes_public_app_when_external_host_is_closed() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = RuntimeHost(_ExternalApplicationRuntime(app)) + host.close() + + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + model_id="fake-runner", + ).run() + + assert result.status == "rejected" + assert result.reason == "busy" + assert result.error is None + assert app.closed + + def test_runner_closes_public_app_when_context_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -257,6 +275,27 @@ def test_runner_closes_public_app_when_context_cleanup_fails() -> None: assert app.closed +def test_runner_preserves_failed_result_without_error_when_cleanup_fails() -> None: + app = _RunnerFakeApplication(total_steps=1) + + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="failed-result", + driver=_FailedResultDriver(), + ), + model_id="fake-runner", + ).run() + + assert result.status == "failed" + assert result.reason == "driver returned failed" + assert result.error is None + assert app.closed + + def test_runner_preserves_primary_error_when_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1, fail_init=True) @@ -309,6 +348,28 @@ async def test_runner_run_async_delegates_to_async_session_helper() -> None: assert io_handler.emitted_steps == [0] +@pytest.mark.asyncio +async def test_runner_run_async_closes_public_app_when_external_host_is_closed() -> ( + None +): + app = _RunnerFakeApplication(total_steps=1) + host = RuntimeHost(_ExternalApplicationRuntime(app)) + host.close() + + result = await Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + + assert result.status == "rejected" + assert result.reason == "busy" + assert result.error is None + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_closes_public_app_when_context_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -325,6 +386,30 @@ async def test_runner_run_async_closes_public_app_when_context_cleanup_fails() - assert app.closed +@pytest.mark.asyncio +async def test_runner_run_async_preserves_failed_result_without_error_when_cleanup_fails() -> ( + None +): + app = _RunnerFakeApplication(total_steps=1) + + result = await Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="async-failed-result", + driver=_AsyncFailedResultDriver(), + ), + model_id="fake-runner", + ).run_async() + + assert result.status == "failed" + assert result.reason == "driver returned failed" + assert result.error is None + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_preserves_primary_error_when_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1, fail_init=True) @@ -872,6 +957,24 @@ def close(self) -> Sequence[OutputArtifact]: return () +class _FailedResultDriver: + def run_one_session( + self, + *, + host: Any, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + host.call(provider.close) + return session_edges.close_result( + status="failed", + reason="driver returned failed", + error=None, + ) + + @dataclass(slots=True) class _AsyncRecordingRunMode: io_handler: _RecordingIOHandler @@ -879,7 +982,7 @@ class _AsyncRecordingRunMode: capabilities: RunModeCapabilities = field( default_factory=lambda: RunModeCapabilities(supports_artifacts=True) ) - driver: "_AsyncBatchDriver" = field(default_factory=lambda: _AsyncBatchDriver()) + driver: Any = field(default_factory=lambda: _AsyncBatchDriver()) def validate_run(self, *, spec: DemoSpec, adapter: Any) -> None: del spec, adapter @@ -928,7 +1031,7 @@ def create_session_edges( cleanup_tasks=context.cleanup_tasks, ) - def select_driver(self) -> "_AsyncBatchDriver": + def select_driver(self) -> Any: return self.driver @@ -953,6 +1056,24 @@ async def run_one_session( ) +class _AsyncFailedResultDriver: + async def run_one_session( + self, + *, + host: Any, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + await host.call_async(provider.close) + return session_edges.close_result( + status="failed", + reason="driver returned failed", + error=None, + ) + + class _AsyncRunModeInputSource: is_finite = False is_deterministic = False From fadac4ce3624426b2e157d251464cb6aed78714a Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:50:30 +0000 Subject: [PATCH 12/33] Add named pull input state API Add InputName and central input state decoders for public demo IO handlers. ReplayIOHandler now exposes get_user_input_state as a deterministic view over the current UserInputWindow, including keyboard state, snapshot-backed state, and legacy key_* compatibility. Add fake CPU tests for provider access, pull/window parity, and named state lookup through the public API. --- flashdreams/flashdreams/demo/__init__.py | 18 ++ flashdreams/flashdreams/demo/application.py | 4 +- flashdreams/flashdreams/demo/inputs.py | 267 ++++++++++++++++++ flashdreams/flashdreams/demo/io.py | 30 +- .../tests/test_demo_application_api.py | 213 ++++++++++++++ 5 files changed, 526 insertions(+), 6 deletions(-) create mode 100644 flashdreams/flashdreams/demo/inputs.py diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 7409dea0a..4edad1571 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -16,6 +16,16 @@ IOutputSink, RuntimeOutputSinkFrameAdapter, ) +from flashdreams.demo.inputs import ( + InputName, + InputStateDecoder, + InputStateDecoderRegistry, + KeyboardInputState, + KeyboardInputStateDecoder, + SnapshotInputStateDecoder, + create_default_input_state_decoder_registry, + input_state_from_window, +) from flashdreams.demo.io import ( CallbackIOHandlerServer, IOHandlerServer, @@ -40,14 +50,22 @@ "IOutputSink", "IOHandler", "IOHandlerServer", + "InputName", + "InputStateDecoder", + "InputStateDecoderRegistry", "InferenceSessionApplicationAdapter", + "KeyboardInputState", + "KeyboardInputStateDecoder", "NativeWindowIOHandler", "ReplayIOHandler", "RuntimeOutputSinkFrameAdapter", "Runner", + "SnapshotInputStateDecoder", "WebRTCIOHandlerServer", "create_native_window_io_handler", + "create_default_input_state_decoder_registry", "create_replay_io_handler", "create_webrtc_io_handler", + "input_state_from_window", "run_replay_application", ] diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py index 4e86957ff..33b23fae4 100644 --- a/flashdreams/flashdreams/demo/application.py +++ b/flashdreams/flashdreams/demo/application.py @@ -33,6 +33,8 @@ step_requirements_from_request, ) +from .inputs import InputName + @runtime_checkable class ApplicationSession(Protocol): @@ -100,7 +102,7 @@ def next_window(self, requirements: StepRequirements) -> UserInputWindow: """ ... - def get_user_input_state(self, modality: str, name: str) -> Any: + def get_user_input_state(self, modality: str, name: InputName | str) -> Any: """Return the current named input state for interactive applications.""" ... diff --git a/flashdreams/flashdreams/demo/inputs.py b/flashdreams/flashdreams/demo/inputs.py new file mode 100644 index 000000000..3c9e4c51f --- /dev/null +++ b/flashdreams/flashdreams/demo/inputs.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Named pull-based input state for public demo IO handlers.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Protocol + +from flashdreams.runtime.demo.session_inputs import UserInputWindow +from flashdreams.runtime.inputs import UserInputEvent, UserInputs +from flashdreams.runtime.keyboard import ( + DEFAULT_SUPPORTED_KEYS, + DRIVING_SUPPORTED_KEYS, + KeyboardState, + WSAD_SUPPORTED_KEYS, + normalize_key, +) + + +class InputName(str, Enum): + """Closed public names for pull-based input state queries.""" + + KEYBOARD = "keyboard" + MOUSE_POSITION = "mouse_position" + MOUSE_BUTTON = "mouse_button" + HEAD_POSITION = "head_position" + HAND_POSITION = "hand_position" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class KeyboardInputState: + """Keyboard state derived from one deterministic user-input window.""" + + pressed_keys: frozenset[str] = field(default_factory=frozenset) + effective_keys: frozenset[str] = field(default_factory=frozenset) + + def is_pressed(self, key: str) -> bool: + """Return whether ``key`` is pressed in this state view.""" + return normalize_key(key) in self.pressed_keys + + +class InputStateDecoder(Protocol): + """Decode one named input state from the current user-input window.""" + + name: InputName + + def state_from_window( + self, + *, + modality: str, + window: UserInputWindow, + ) -> Any: + """Return the named input state for ``window``.""" + ... + + +class InputStateDecoderRegistry: + """Registry of named input decoders used by public IO handlers.""" + + def __init__(self, decoders: Sequence[InputStateDecoder] = ()) -> None: + self._decoders: dict[InputName, InputStateDecoder] = {} + for decoder in decoders: + self.register(decoder) + + def register(self, decoder: InputStateDecoder) -> None: + """Register one named input decoder.""" + decoder_name = getattr(decoder, "name", None) + if not isinstance(decoder_name, InputName) or not callable( + getattr(decoder, "state_from_window", None) + ): + raise TypeError("decoder must implement the InputStateDecoder protocol.") + if decoder_name in self._decoders: + raise ValueError(f"Input decoder {decoder_name.value!r} is already set.") + self._decoders[decoder_name] = decoder + + @property + def decoders(self) -> Mapping[InputName, InputStateDecoder]: + """Return registered decoders by name.""" + return dict(self._decoders) + + def state_from_window( + self, + *, + modality: str, + name: InputName | str, + window: UserInputWindow, + ) -> Any: + """Return one named state view over ``window``.""" + key_name = _legacy_key_name(name) + if key_name is not None: + keyboard = self.state_from_window( + modality=modality, + name=InputName.KEYBOARD, + window=window, + ) + if isinstance(keyboard, KeyboardInputState): + return keyboard.is_pressed(key_name) + return False + + input_name = _coerce_input_name(name) + if input_name is None: + return None + decoder = self._decoders.get(input_name) + if decoder is None: + return None + return decoder.state_from_window(modality=modality, window=window) + + +@dataclass(frozen=True, slots=True) +class KeyboardInputStateDecoder: + """Decode keyboard edge events into a pullable level-triggered state.""" + + name: InputName = InputName.KEYBOARD + + def state_from_window( + self, + *, + modality: str, + window: UserInputWindow, + ) -> KeyboardInputState | None: + if modality.strip().lower() not in {"", "keyboard"}: + return None + + inputs = window.inputs + initial_keys = _snapshot_pressed_keys(inputs.snapshot) + state = KeyboardState( + pressed_keys=set(initial_keys), + supported_keys=_supported_keys(inputs, initial_keys), + ) + for event in inputs.events: + keyboard_event = _keyboard_event_name(event) + if keyboard_event is None: + continue + key = event.payload.get("key") + if isinstance(key, str): + state.apply_event(event=keyboard_event, key=key) + return KeyboardInputState( + pressed_keys=state.snapshot(), + effective_keys=state.resolved_effective_keys(), + ) + + +@dataclass(frozen=True, slots=True) +class SnapshotInputStateDecoder: + """Return named non-keyboard values from the current window snapshot.""" + + name: InputName + + def state_from_window( + self, + *, + modality: str, + window: UserInputWindow, + ) -> Any: + del modality + snapshot = window.inputs.snapshot + if self.name.value in snapshot: + return snapshot[self.name.value] + nested = snapshot.get(self.name.value.split("_", maxsplit=1)[0]) + if isinstance(nested, Mapping): + return nested.get(self.name.value) + return None + + +def create_default_input_state_decoder_registry() -> InputStateDecoderRegistry: + """Create the standard named input decoder registry.""" + return InputStateDecoderRegistry( + ( + KeyboardInputStateDecoder(), + SnapshotInputStateDecoder(InputName.MOUSE_POSITION), + SnapshotInputStateDecoder(InputName.MOUSE_BUTTON), + SnapshotInputStateDecoder(InputName.HEAD_POSITION), + SnapshotInputStateDecoder(InputName.HAND_POSITION), + ) + ) + + +def input_state_from_window( + window: UserInputWindow, + *, + modality: str, + name: InputName | str, +) -> Any: + """Decode one named input state from ``window`` using default decoders.""" + return create_default_input_state_decoder_registry().state_from_window( + modality=modality, + name=name, + window=window, + ) + + +def _coerce_input_name(name: InputName | str) -> InputName | None: + if isinstance(name, InputName): + return name + try: + return InputName(name) + except ValueError: + return None + + +def _legacy_key_name(name: InputName | str) -> str | None: + if not isinstance(name, str): + return None + if not name.startswith("key_"): + return None + return name.removeprefix("key_") + + +def _snapshot_pressed_keys(snapshot: Mapping[str, Any]) -> frozenset[str]: + value = snapshot.get("pressed_keys") + if value is None: + keyboard = snapshot.get("keyboard") + if isinstance(keyboard, Mapping): + value = keyboard.get("pressed_keys") + if isinstance(value, str): + return frozenset({normalize_key(value)}) + if isinstance(value, Sequence): + return frozenset( + normalize_key(key) + for key in value + if isinstance(key, str) and key.strip() + ) + return frozenset() + + +def _supported_keys( + inputs: UserInputs, + initial_keys: frozenset[str], +) -> frozenset[str]: + event_keys = { + normalize_key(key) + for event in inputs.events + for key in (event.payload.get("key"),) + if isinstance(key, str) and key.strip() + } + return frozenset( + set(DEFAULT_SUPPORTED_KEYS) + | set(DRIVING_SUPPORTED_KEYS) + | set(WSAD_SUPPORTED_KEYS) + | set(initial_keys) + | event_keys + ) + + +def _keyboard_event_name(event: UserInputEvent) -> str | None: + normalized = event.event_type.strip().lower().replace(".", "_") + if normalized in {"key_down", "keyboard_keydown", "keydown"}: + return "keydown" + if normalized in {"key_up", "keyboard_keyup", "keyup"}: + return "keyup" + return None + + +__all__ = [ + "InputName", + "InputStateDecoder", + "InputStateDecoderRegistry", + "KeyboardInputState", + "KeyboardInputStateDecoder", + "SnapshotInputStateDecoder", + "create_default_input_state_decoder_registry", + "input_state_from_window", +] diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py index 79fd650f9..9c9f98352 100644 --- a/flashdreams/flashdreams/demo/io.py +++ b/flashdreams/flashdreams/demo/io.py @@ -35,6 +35,11 @@ from flashdreams.runtime.types import StepRequirements, StepResult from .application import FrameOutputSink, IOHandler +from .inputs import ( + InputName, + InputStateDecoderRegistry, + create_default_input_state_decoder_registry, +) RunSessionCallback = Callable[[IOHandler], RunResult] ServeCallback = Callable[[], object] @@ -59,8 +64,16 @@ class ReplayIOHandler: is_finite: bool = True is_deterministic: bool = True user_input_schema: UserInputSchema = field(default_factory=UserInputSchema) + input_state_decoders: InputStateDecoderRegistry = field( + default_factory=create_default_input_state_decoder_registry + ) _input_source: "_ReplayIOInputSource" = field(init=False, repr=False) _output_sink: OutputSink | FrameOutputSink = field(init=False, repr=False) + _current_window: UserInputWindow | None = field( + default=None, + init=False, + repr=False, + ) _opened_session_info: SessionInfo | None = field( default=None, init=False, @@ -93,11 +106,18 @@ def open(self, session_info: SessionInfo) -> None: open_output(session_info) def next_window(self, requirements: StepRequirements) -> UserInputWindow: - return self._input_source.next_window(requirements) - - def get_user_input_state(self, modality: str, name: str) -> Any: - del modality, name - return None + window = self._input_source.next_window(requirements) + self._current_window = window + return window + + def get_user_input_state(self, modality: str, name: InputName | str) -> Any: + if self._current_window is None: + return None + return self.input_state_decoders.state_from_window( + modality=modality, + name=name, + window=self._current_window, + ) def begin_generation(self, generation: int) -> None: self._generation = generation diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index a26ccc32d..2003b422f 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -19,8 +19,10 @@ DemoApplication, FrameOutputSink, InferenceSessionApplicationAdapter, + InputName, IOHandler, IOHandlerServer, + KeyboardInputState, ReplayIOHandler, Runner, RuntimeOutputSinkFrameAdapter, @@ -28,6 +30,7 @@ create_native_window_io_handler, create_replay_io_handler, create_webrtc_io_handler, + input_state_from_window, ) from flashdreams.runtime import ( CanonicalInputSchema, @@ -42,6 +45,7 @@ StepRequest, StepResult, TimeWindow, + UserInputEvent, UserInputs, UserInputSchema, ) @@ -177,6 +181,72 @@ def test_io_handler_protocol_keeps_input_conversion_outside_io() -> None: assert handler.get_user_input_state("keyboard", "key_w") is False +def test_replay_io_handler_pulls_keyboard_state_from_current_window() -> None: + inputs = UserInputs( + events=( + _key_event("key_down", "w", 0.1), + _key_event("keydown", "ArrowLeft", 0.2), + _key_event("key_up", "w", 0.8), + ) + ) + handler = create_replay_io_handler(replay_log=inputs) + + assert handler.get_user_input_state("keyboard", InputName.KEYBOARD) is None + window = handler.next_window(StepRequirements(step_index=0, input_frame_count=1)) + pulled = handler.get_user_input_state("keyboard", InputName.KEYBOARD) + direct = input_state_from_window( + window, + modality="keyboard", + name=InputName.KEYBOARD, + ) + + assert pulled == direct + assert isinstance(pulled, KeyboardInputState) + assert pulled.pressed_keys == frozenset({"a"}) + assert pulled.effective_keys == frozenset({"a"}) + assert handler.get_user_input_state("keyboard", "key_a") is True + assert handler.get_user_input_state("keyboard", "key_w") is False + + +def test_replay_io_handler_pulls_snapshot_backed_named_state() -> None: + handler = create_replay_io_handler( + replay_log=UserInputs( + snapshot={ + "mouse": {"mouse_position": (10, 20)}, + "hand_position": (3, 4, 5), + } + ) + ) + + handler.next_window(StepRequirements(step_index=0)) + + assert handler.get_user_input_state("mouse", InputName.MOUSE_POSITION) == (10, 20) + assert handler.get_user_input_state("hand", InputName.HAND_POSITION) == (3, 4, 5) + assert handler.get_user_input_state("unknown", "unknown_state") is None + + +def test_provider_can_pull_keyboard_state_through_replay_io_handler() -> None: + inputs = UserInputs(events=(_key_event("key_down", "w", 0.1),)) + io_handler = create_replay_io_handler(replay_log=inputs) + adapter = _PullStateDemoAdapter(io_handler) + app = DemoAdapterApplication( + adapter=adapter, + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + config=InferenceConfig(model_id="fake-demo"), + ), + ) + + result = Runner(io_handler=io_handler, app=app).run() + + assert result.status == "completed" + assert adapter.runtime.session.step_inputs == [ + {"effective_keys": frozenset({"w"}), "pressed_keys": frozenset({"w"})} + ] + + def test_runtime_output_sink_frame_adapter_satisfies_frame_output_sink() -> None: output = NullOutputSink(store_results=True) output.open(SessionInfo()) @@ -547,6 +617,14 @@ def test_demo_application_server_selection_does_not_build_replay_app() -> None: assert not app.replay_adapter_called +def _key_event(event_type: str, key: str, timestamp_s: float) -> UserInputEvent: + return UserInputEvent( + timestamp_s=timestamp_s, + event_type=event_type, + payload={"key": key}, + ) + + class _ReplayDemoApplication(DemoApplication): def __init__(self) -> None: self.adapter = _FakeDemoAdapter() @@ -656,6 +734,141 @@ def close(self) -> None: self.closed = True +class _PullStateDemoAdapter: + model_id = "fake-demo" + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name="prompt"),), + ) + canonical_input_schema = CanonicalInputSchema() + + def __init__(self, io_handler: IOHandler) -> None: + self.io_handler = io_handler + self.runtime = _PullStateRuntime() + self.runtimes: list[_PullStateRuntime] = [] + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("null",) + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + assert config.model_id == self.model_id + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + self.runtime = _PullStateRuntime() + self.runtimes.append(self.runtime) + return self.runtime + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + assert spec.model_id == self.model_id + return PreparedScenario( + initial_inputs=InferenceInput(global_conditioning={"prompt": "demo"}), + user_inputs=UserInputs(), + source_schema=UserInputSchema(), + ) + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: PreparedScenario, + ) -> "_PullStateModelInputProvider": + del scenario + assert spec.model_id == self.model_id + return _PullStateModelInputProvider(self.io_handler) + + +class _PullStateModelInputProvider: + capabilities = ProviderCapabilities( + supports_recorded_input=True, + deterministic_given_inputs=True, + ) + + def __init__(self, io_handler: IOHandler) -> None: + self.io_handler = io_handler + self.closed = False + + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput(global_conditioning={"prompt": "demo"}) + + def prepare_step( + self, + *, + request: StepRequirements, + user_window: UserInputWindow, + ) -> PreparedStep: + del request, user_window + state = self.io_handler.get_user_input_state("keyboard", InputName.KEYBOARD) + if not isinstance(state, KeyboardInputState): + raise RuntimeError("expected keyboard state") + return PreparedStep( + inference_input=InferenceInput( + step={ + "effective_keys": state.effective_keys, + "pressed_keys": state.pressed_keys, + } + ) + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + self.closed = True + + +class _PullStateRuntime: + def __init__(self) -> None: + self.session = _PullStateSession() + self.closed = False + + def start_session(self, inputs: InferenceInput) -> "_PullStateSession": + assert inputs.global_conditioning["prompt"] == "demo" + return self.session + + def close(self) -> None: + self.closed = True + + +class _PullStateSession: + def __init__(self) -> None: + self.step_inputs: list[dict[str, frozenset[str]]] = [] + self.closed = False + + def session_info(self) -> SessionInfo: + return SessionInfo(output_layout="pull-state") + + def next_step_requirements(self) -> StepRequirements | None: + if self.step_inputs: + return None + return StepRequirements(step_index=0) + + def next_step_request(self) -> StepRequest | None: + if self.step_inputs: + return None + return StepRequest(step_index=0) + + def step(self, inputs: InferenceInput) -> StepResult: + self.step_inputs.append( + { + "effective_keys": inputs.step["effective_keys"], + "pressed_keys": inputs.step["pressed_keys"], + } + ) + return StepResult(step_index=0, output="pull-state", frame_count=1) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + self.step_inputs.clear() + + def close(self) -> None: + self.closed = True + + class _FakeDemoAdapter: model_id = "fake-demo" inference_input_schema = InferenceInputSchema( From 0ad73e51691787c7c14659f58847241e1858de90 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 08:59:02 +0000 Subject: [PATCH 13/33] Shield async public runner cleanup --- flashdreams/flashdreams/demo/inputs.py | 6 +- flashdreams/flashdreams/demo/runner.py | 57 +++++++++++++--- .../tests/test_demo_application_api.py | 65 +++++++++++++++++++ 3 files changed, 116 insertions(+), 12 deletions(-) diff --git a/flashdreams/flashdreams/demo/inputs.py b/flashdreams/flashdreams/demo/inputs.py index 3c9e4c51f..c32d50302 100644 --- a/flashdreams/flashdreams/demo/inputs.py +++ b/flashdreams/flashdreams/demo/inputs.py @@ -15,8 +15,8 @@ from flashdreams.runtime.keyboard import ( DEFAULT_SUPPORTED_KEYS, DRIVING_SUPPORTED_KEYS, - KeyboardState, WSAD_SUPPORTED_KEYS, + KeyboardState, normalize_key, ) @@ -220,9 +220,7 @@ def _snapshot_pressed_keys(snapshot: Mapping[str, Any]) -> frozenset[str]: return frozenset({normalize_key(value)}) if isinstance(value, Sequence): return frozenset( - normalize_key(key) - for key in value - if isinstance(key, str) and key.strip() + normalize_key(key) for key in value if isinstance(key, str) and key.strip() ) return frozenset() diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 0aac516f2..67468decf 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -7,7 +7,7 @@ import asyncio import inspect -from collections.abc import Sequence +from collections.abc import Coroutine, Sequence from dataclasses import dataclass, field from typing import Any, cast @@ -16,6 +16,7 @@ from flashdreams.runtime.demo.drivers import ( run_demo_session, run_demo_session_async, + uncancel_current_task, ) from flashdreams.runtime.demo.host import ModelWarmupPlan, RuntimeHost from flashdreams.runtime.demo.pipeline import StepPipeline @@ -108,13 +109,15 @@ async def run_async(self) -> RunResult: primary_error = exc raise finally: - await _close_runner_resources_async( - context=context, - host=host, - app=self.app, - owns_host=owns_host, - run_result=result, - primary_error=primary_error, + await _await_runner_cleanup( + _close_runner_resources_async( + context=context, + host=host, + app=self.app, + owns_host=owns_host, + run_result=result, + primary_error=primary_error, + ) ) def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: @@ -393,6 +396,33 @@ async def _close_runner_resources_async( _raise_first_cleanup_error(errors) +async def _await_runner_cleanup(cleanup: Coroutine[Any, Any, None]) -> None: + cleanup_task = asyncio.create_task(cleanup) + was_cancelled = False + cleanup_error: BaseException | None = None + + while True: + try: + await asyncio.shield(cleanup_task) + break + except asyncio.CancelledError as exc: + if cleanup_task.done(): + cleanup_error = exc + break + was_cancelled = True + uncancel_current_task() + except Exception as exc: + cleanup_error = exc + break + + if was_cancelled: + cancellation = asyncio.CancelledError("cancelled during runner cleanup") + _record_cancelled_cleanup_note(cancellation, cleanup_error) + raise cancellation from None + if cleanup_error is not None: + raise cleanup_error + + def _record_cleanup_error( errors: list[Exception], cleanup: Any, @@ -469,6 +499,17 @@ def _record_cleanup_notes( add_note(f"Additional cleanup error: {extra!r}") +def _record_cancelled_cleanup_note( + cancellation: asyncio.CancelledError, + cleanup_error: BaseException | None, +) -> None: + if cleanup_error is None: + return + add_note = getattr(cancellation, "add_note", None) + if callable(add_note): + add_note(f"Cleanup failed: {cleanup_error!r}") + + def _application_adapter(app: Application) -> DemoAdapter | None: if isinstance(app, DemoAdapterApplication): return app.adapter diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 2003b422f..07b39a6ac 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +import asyncio import threading from collections.abc import Sequence from dataclasses import dataclass, field @@ -11,6 +12,7 @@ import pytest +import flashdreams.demo.runner as demo_runner from flashdreams.demo import ( Application, ApplicationSession, @@ -516,6 +518,50 @@ async def test_runner_run_async_preserves_failed_result_when_cleanup_fails() -> assert app.closed +@pytest.mark.asyncio +async def test_runner_run_async_finishes_cleanup_when_cancelled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + close_started = threading.Event() + release_close = threading.Event() + created_hosts: list[RuntimeHost] = [] + + class _ObservedRuntimeHost(RuntimeHost): + def __init__(self, runtime: InferenceRuntime) -> None: + super().__init__(runtime) + created_hosts.append(self) + + monkeypatch.setattr(demo_runner, "RuntimeHost", _ObservedRuntimeHost) + app = _BlockingCloseRunnerFakeApplication( + total_steps=1, + close_started=close_started, + release_close=release_close, + ) + runner_task = asyncio.create_task( + Runner( + io_handler=_RecordingIOHandler(), + app=app, + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + ) + + try: + assert await asyncio.to_thread(close_started.wait, 2.0) + runner_task.cancel() + await asyncio.sleep(0) + assert not runner_task.done() + finally: + release_close.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(runner_task, timeout=2.0) + + assert app.closed + assert len(created_hosts) == 1 + assert created_hosts[0].is_closed + + def test_replay_io_factory_runs_through_public_runner() -> None: output_sink = NullOutputSink(store_results=True) io_handler = create_replay_io_handler(output_sink=output_sink) @@ -1046,6 +1092,25 @@ def close(self) -> None: self.closed = True +class _BlockingCloseRunnerFakeApplication(_RunnerFakeApplication): + def __init__( + self, + *, + total_steps: int, + close_started: threading.Event, + release_close: threading.Event, + ) -> None: + super().__init__(total_steps=total_steps) + self._close_started = close_started + self._release_close = release_close + + def close(self) -> None: + self._close_started.set() + if not self._release_close.wait(2.0): + raise RuntimeError("fake close was not released") + super().close() + + class _ExternalApplicationRuntime: def __init__(self, app: _RunnerFakeApplication) -> None: self._app = app From 485a9bd566c4565a0728195ff3d85c8a6362f230 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 09:11:05 +0000 Subject: [PATCH 14/33] Add output comparison sink --- flashdreams/flashdreams/demo/__init__.py | 10 + flashdreams/flashdreams/demo/io.py | 82 +++++-- .../flashdreams/runtime/demo/__init__.py | 6 + .../flashdreams/runtime/demo/outputs.py | 201 ++++++++++++++++++ .../tests/test_demo_application_api.py | 42 ++++ .../tests/test_demo_runtime_output_sinks.py | 122 +++++++++++ 6 files changed, 449 insertions(+), 14 deletions(-) diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 4edad1571..248a4dd68 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -37,13 +37,23 @@ create_webrtc_io_handler, ) from flashdreams.demo.runner import Runner +from flashdreams.runtime.demo.outputs import ( + BenchmarkStatsOutputSink, + ComparisonOutputMismatchError, + ComparisonOutputSink, + FileOutputSink, +) __all__ = [ "Application", "ApplicationSession", + "BenchmarkStatsOutputSink", "CallbackIOHandlerServer", + "ComparisonOutputMismatchError", + "ComparisonOutputSink", "DemoAdapterApplication", "DemoApplication", + "FileOutputSink", "FrameOutputSink", "IApplication", "IApplicationSession", diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py index 9c9f98352..9f856f7f5 100644 --- a/flashdreams/flashdreams/demo/io.py +++ b/flashdreams/flashdreams/demo/io.py @@ -12,6 +12,7 @@ from flashdreams.runtime.demo.drivers import BatchSessionDriver from flashdreams.runtime.demo.host import ModelWarmupPlan, RuntimeHost from flashdreams.runtime.demo.outputs import ( + CompositeOutputSinkError, NullOutputSink, OutputDecision, OutputSink, @@ -101,9 +102,18 @@ def run_mode(self) -> "IOHandlerRunMode": def open(self, session_info: SessionInfo) -> None: self._opened_session_info = session_info - open_output = getattr(self._output_sink, "open", None) - if callable(open_output): - open_output(session_info) + opened: list[object] = [] + errors: list[BaseException] = [] + for tail in _output_tails(self._output_sink, self.metric_output_sink): + try: + if _open_optional_output(tail, session_info): + opened.append(tail) + except Exception as exc: + errors.append(exc) + if errors: + for tail in reversed(opened): + _close_optional_output(tail, artifacts=[], errors=errors) + raise CompositeOutputSinkError("open", errors) def next_window(self, requirements: StepRequirements) -> UserInputWindow: window = self._input_source.next_window(requirements) @@ -121,9 +131,10 @@ def get_user_input_state(self, modality: str, name: InputName | str) -> Any: def begin_generation(self, generation: int) -> None: self._generation = generation - begin_generation = getattr(self._output_sink, "begin_generation", None) - if callable(begin_generation): - begin_generation(generation) + for tail in _output_tails(self._output_sink, self.metric_output_sink): + begin_generation = getattr(tail, "begin_generation", None) + if callable(begin_generation): + begin_generation(generation) def emit_chunk(self, result: StepResult) -> OutputDecision: write = getattr(self._output_sink, "write", None) @@ -139,7 +150,10 @@ def emit_chunk(self, result: StepResult) -> OutputDecision: handle_output = getattr(self._output_sink, "handle_output") handle_output(timestamp_s, result) decision = OutputDecision() - if self.metric_output_sink is not None: + if ( + self.metric_output_sink is not None + and self.metric_output_sink is not self._output_sink + ): self.metric_output_sink.handle_output(_result_timestamp_s(result), result) if decision.should_stop: self._should_exit = True @@ -151,13 +165,13 @@ def should_exit(self) -> bool: def close(self) -> Sequence[OutputArtifact]: self._closed = True self._should_exit = True - close = getattr(self._output_sink, "close", None) - if callable(close): - artifacts = close() - if artifacts is None: - return () - return tuple(artifacts) - return () + artifacts: list[OutputArtifact] = [] + errors: list[BaseException] = [] + for tail in _output_tails(self._output_sink, self.metric_output_sink): + _close_optional_output(tail, artifacts=artifacts, errors=errors) + if errors: + raise CompositeOutputSinkError("close", errors) + return tuple(artifacts) @dataclass(slots=True) @@ -329,6 +343,46 @@ def close(self) -> Sequence[OutputArtifact]: return self.io_handler.close() +def _output_tails(*tails: object | None) -> tuple[object, ...]: + distinct: list[object] = [] + seen: set[int] = set() + for tail in tails: + if tail is None: + continue + tail_id = id(tail) + if tail_id in seen: + continue + seen.add(tail_id) + distinct.append(tail) + return tuple(distinct) + + +def _open_optional_output(tail: object, session_info: SessionInfo) -> bool: + open_output = getattr(tail, "open", None) + if callable(open_output): + open_output(session_info) + return True + return False + + +def _close_optional_output( + tail: object, + *, + artifacts: list[OutputArtifact], + errors: list[BaseException], +) -> None: + close = getattr(tail, "close", None) + if not callable(close): + return + try: + closed_artifacts = close() + except Exception as exc: + errors.append(exc) + return + if closed_artifacts is not None: + artifacts.extend(closed_artifacts) + + def create_replay_io_handler( replay_log: UserInputs | None = None, output_sink: OutputSink | FrameOutputSink | None = None, diff --git a/flashdreams/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index b3d6598dd..515545478 100644 --- a/flashdreams/flashdreams/runtime/demo/__init__.py +++ b/flashdreams/flashdreams/runtime/demo/__init__.py @@ -25,8 +25,11 @@ ) from flashdreams.runtime.demo.outputs import ( BenchmarkStatsOutputSink, + ComparisonOutputMismatchError, + ComparisonOutputSink, CompositeOutputSink, CompositeOutputSinkError, + FileOutputSink, Mp4OutputSink, NullOutputSink, OutputDecision, @@ -133,6 +136,8 @@ "BenchmarkStatsOutputSink", "BenchmarkRunMode", "BenchmarkErrorPolicy", + "ComparisonOutputMismatchError", + "ComparisonOutputSink", "InMemorySessionMetricsRecorder", "InputSource", "CatchUpDecision", @@ -140,6 +145,7 @@ "CompositeOutputSink", "CompositeOutputSinkError", "DeterministicClock", + "FileOutputSink", "MetricsSnapshot", "ModelWarmupAdapter", "ModelWarmupPlan", diff --git a/flashdreams/flashdreams/runtime/demo/outputs.py b/flashdreams/flashdreams/runtime/demo/outputs.py index 59f53d7ed..ebebe98e7 100644 --- a/flashdreams/flashdreams/runtime/demo/outputs.py +++ b/flashdreams/flashdreams/runtime/demo/outputs.py @@ -12,6 +12,8 @@ from pathlib import Path from typing import Any, Literal, Protocol, runtime_checkable +import torch + from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner_io import ( DEFAULT_RUNNER_INSTALL_HINT, @@ -102,6 +104,10 @@ def __init__(self, operation: str, errors: Sequence[BaseException]) -> None: ) +class ComparisonOutputMismatchError(AssertionError): + """Raised when generated output differs from an expected CI baseline.""" + + @dataclass(slots=True) class NullOutputSink: """Output sink for headless runs and fake-model vertical-slice tests.""" @@ -135,6 +141,10 @@ def write(self, result: StepResult) -> OutputDecision: self.results.append(_result_record(result)) return OutputDecision() + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.write(chunk) + def close(self) -> Sequence[OutputArtifact]: self.closed = True return () @@ -201,6 +211,10 @@ def write(self, result: StepResult) -> OutputDecision: self._collector.add(result) return OutputDecision() + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.write(chunk) + def close(self) -> Sequence[OutputArtifact]: if self._artifacts is not None: return self._artifacts @@ -244,6 +258,9 @@ def close(self) -> Sequence[OutputArtifact]: return self._artifacts +FileOutputSink = Mp4OutputSink + + @dataclass(slots=True) class BenchmarkStatsOutputSink: """Structured benchmark metrics artifact sink for shared demo runs.""" @@ -289,6 +306,10 @@ def write(self, result: StepResult) -> OutputDecision: self._samples.extend(samples) return OutputDecision() + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.write(chunk) + def close(self) -> Sequence[OutputArtifact]: if self._artifacts is not None: return self._artifacts @@ -324,6 +345,84 @@ def close(self) -> Sequence[OutputArtifact]: return self._artifacts +@dataclass(slots=True) +class ComparisonOutputSink: + """Compare generated outputs against an expected deterministic sequence.""" + + expected_results: Sequence[StepResult] + compare_output: bool = True + compare_output_window: bool = True + compare_metrics: bool = False + compare_metadata: bool = False + rtol: float = 0.0 + atol: float = 0.0 + produces_artifacts: bool = False + _opened: bool = field(default=False, init=False, repr=False) + _closed: bool = field(default=True, init=False, repr=False) + _position: int = field(default=0, init=False, repr=False) + + def __post_init__(self) -> None: + self.expected_results = tuple(self.expected_results) + if self.rtol < 0: + raise ValueError("ComparisonOutputSink.rtol must be >= 0.") + if self.atol < 0: + raise ValueError("ComparisonOutputSink.atol must be >= 0.") + + def open(self, session_info: SessionInfo) -> None: + del session_info + self._position = 0 + self._opened = True + self._closed = False + + def begin_generation(self, generation: int) -> None: + if generation < 0: + raise ValueError("generation must be >= 0.") + + def write(self, result: StepResult) -> OutputDecision: + if not self._opened or self._closed: + raise RuntimeError("Cannot write to a closed output sink.") + if self._position >= len(self.expected_results): + raise ComparisonOutputMismatchError( + "Unexpected output at position " + f"{self._position}: step_index={result.step_index}." + ) + + expected = self.expected_results[self._position] + mismatches = _compare_step_result( + expected=expected, + actual=result, + compare_output=self.compare_output, + compare_output_window=self.compare_output_window, + compare_metrics=self.compare_metrics, + compare_metadata=self.compare_metadata, + rtol=self.rtol, + atol=self.atol, + ) + if mismatches: + details = "; ".join(mismatches) + raise ComparisonOutputMismatchError( + f"Output mismatch at position {self._position}: {details}" + ) + + self._position += 1 + return OutputDecision() + + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.write(chunk) + + def close(self) -> Sequence[OutputArtifact]: + self._opened = False + self._closed = True + if self._position != len(self.expected_results): + missing = len(self.expected_results) - self._position + raise ComparisonOutputMismatchError( + f"Missing {missing} expected output(s); " + f"received {self._position} of {len(self.expected_results)}." + ) + return () + + @dataclass(slots=True) class CompositeOutputSink: """Fan out generated outputs to multiple sinks and return all artifacts.""" @@ -382,6 +481,10 @@ def write(self, result: StepResult) -> OutputDecision: raise RuntimeError("Cannot write to a closed output sink.") return _combine_output_decisions(sink.write(result) for sink in self.sinks) + def handle_output(self, timestamp_s: float, chunk: StepResult) -> None: + del timestamp_s + self.write(chunk) + def close(self) -> Sequence[OutputArtifact]: if self._artifacts is not None: return self._artifacts @@ -441,6 +544,101 @@ def build_benchmark_output_sink( ) +def _compare_step_result( + *, + expected: StepResult, + actual: StepResult, + compare_output: bool, + compare_output_window: bool, + compare_metrics: bool, + compare_metadata: bool, + rtol: float, + atol: float, +) -> tuple[str, ...]: + mismatches: list[str] = [] + _compare_field(mismatches, "step_index", expected.step_index, actual.step_index) + _compare_field(mismatches, "frame_count", expected.frame_count, actual.frame_count) + _compare_field(mismatches, "layout", expected.layout, actual.layout) + if compare_output_window: + _compare_field( + mismatches, + "output_window", + expected.output_window, + actual.output_window, + ) + if compare_metrics: + _compare_field(mismatches, "metrics", expected.metrics, actual.metrics) + if compare_metadata: + _compare_field(mismatches, "metadata", expected.metadata, actual.metadata) + if compare_output: + _compare_output( + mismatches, + expected.output, + actual.output, + rtol=rtol, + atol=atol, + ) + return tuple(mismatches) + + +def _compare_field( + mismatches: list[str], + name: str, + expected: object, + actual: object, +) -> None: + if actual != expected: + mismatches.append(f"{name} expected {expected!r}, got {actual!r}") + + +def _compare_output( + mismatches: list[str], + expected: object, + actual: object, + *, + rtol: float, + atol: float, +) -> None: + if isinstance(expected, torch.Tensor) or isinstance(actual, torch.Tensor): + if not isinstance(expected, torch.Tensor) or not isinstance( + actual, torch.Tensor + ): + mismatches.append( + "output tensor type expected " + f"{type(expected).__name__}, got {type(actual).__name__}" + ) + return + if tuple(actual.shape) != tuple(expected.shape): + mismatches.append( + f"output shape expected {tuple(expected.shape)}, " + f"got {tuple(actual.shape)}" + ) + return + if actual.dtype != expected.dtype: + mismatches.append( + f"output dtype expected {expected.dtype}, got {actual.dtype}" + ) + return + actual_cpu = actual.detach().cpu() + expected_cpu = expected.detach().cpu() + if actual_cpu.is_floating_point() or expected_cpu.is_floating_point(): + matches = torch.allclose( + actual_cpu, + expected_cpu, + rtol=rtol, + atol=atol, + equal_nan=True, + ) + else: + matches = torch.equal(actual_cpu, expected_cpu) + if not matches: + mismatches.append("output tensor values differ") + return + + if actual != expected: + mismatches.append(f"output expected {expected!r}, got {actual!r}") + + def _result_record(result: StepResult) -> Mapping[str, object]: record: dict[str, object] = { "step_index": result.step_index, @@ -622,8 +820,11 @@ def build_output_target( __all__ = [ "BenchmarkStatsOutputSink", + "ComparisonOutputMismatchError", + "ComparisonOutputSink", "CompositeOutputSinkError", "CompositeOutputSink", + "FileOutputSink", "Mp4OutputSink", "NullOutputSink", "OutputDecision", diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 07b39a6ac..06b048b0a 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -8,6 +8,7 @@ import threading from collections.abc import Sequence from dataclasses import dataclass, field +from pathlib import Path from typing import Any import pytest @@ -16,9 +17,12 @@ from flashdreams.demo import ( Application, ApplicationSession, + BenchmarkStatsOutputSink, CallbackIOHandlerServer, + ComparisonOutputSink, DemoAdapterApplication, DemoApplication, + FileOutputSink, FrameOutputSink, InferenceSessionApplicationAdapter, InputName, @@ -80,7 +84,10 @@ def test_public_demo_contracts_are_importable() -> None: assert Application.__name__ == "Application" assert ApplicationSession.__name__ == "ApplicationSession" + assert BenchmarkStatsOutputSink.__name__ == "BenchmarkStatsOutputSink" assert DemoApplication.__name__ == "DemoApplication" + assert ComparisonOutputSink.__name__ == "ComparisonOutputSink" + assert FileOutputSink.__name__ == "Mp4OutputSink" assert IOHandler.__name__ == "IOHandler" assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" @@ -597,6 +604,25 @@ def test_replay_io_factory_should_exit_tracks_output_stop() -> None: assert io_handler.should_exit() +def test_replay_io_factory_can_gate_ci_output_correctness() -> None: + expected = StepResult( + step_index=0, + output="runner-chunk-0", + frame_count=1, + output_window=TimeWindow(start_s=0.0, end_s=1.0), + ) + comparison_tail = ComparisonOutputSink((expected,)) + io_handler = create_replay_io_handler(output_sink=comparison_tail) + + result = Runner( + io_handler=io_handler, + app=_RunnerFakeApplication(total_steps=1), + model_id="fake-replay-factory", + ).run() + + assert result.status == "completed" + + def test_replay_io_factory_forwards_to_metric_tail() -> None: metric_tail = _RecordingFrameOutputSink() io_handler = create_replay_io_handler(metric_output_sink=metric_tail) @@ -611,6 +637,22 @@ def test_replay_io_factory_forwards_to_metric_tail() -> None: assert metric_tail.records == [(0.0, 0)] +def test_replay_io_factory_closes_benchmark_metric_tail(tmp_path: Path) -> None: + stats_path = tmp_path / "replay-stats.json" + metric_tail = BenchmarkStatsOutputSink(output_path=stats_path) + io_handler = create_replay_io_handler(metric_output_sink=metric_tail) + + result = Runner( + io_handler=io_handler, + app=_RunnerFakeApplication(total_steps=1), + model_id="fake-replay-factory", + ).run() + + assert result.status == "completed" + assert tuple(artifact.uri for artifact in result.artifacts) == (str(stats_path),) + assert stats_path.exists() + + def test_native_window_factory_exposes_io_handler_shape() -> None: io_handler = create_native_window_io_handler((1280, 720)) diff --git a/flashdreams/tests/test_demo_runtime_output_sinks.py b/flashdreams/tests/test_demo_runtime_output_sinks.py index 1d24c7752..3a2e377aa 100644 --- a/flashdreams/tests/test_demo_runtime_output_sinks.py +++ b/flashdreams/tests/test_demo_runtime_output_sinks.py @@ -15,9 +15,12 @@ from flashdreams.runtime import OutputArtifact, StepResult, TimeWindow from flashdreams.runtime.demo import ( BenchmarkStatsOutputSink, + ComparisonOutputMismatchError, + ComparisonOutputSink, CompositeOutputSink, CompositeOutputSinkError, DemoSpec, + FileOutputSink, Mp4OutputSink, Mp4OutputSpec, NullOutputSink, @@ -109,6 +112,57 @@ def fake_writer( ] +def test_file_output_sink_exposes_narrow_frame_tail(tmp_path: Path) -> None: + writer_calls: list[dict[str, Any]] = [] + + def fake_writer( + video: torch.Tensor, + path: Path, + *, + fps: int | float, + layout: str, + install_hint: str, + ) -> Path: + del install_hint + writer_calls.append( + { + "shape": tuple(video.shape), + "path": path, + "fps": fps, + "layout": layout, + } + ) + return path + + sink = FileOutputSink( + output_path=tmp_path / "tail.mp4", + fps=8, + writer=fake_writer, + move_to_cpu=False, + ) + sink.open(SessionInfo(output_layout="bvtchw")) + + sink.handle_output( + 4.0, + StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2)), + layout="bvtchw", + ), + ) + artifacts = tuple(sink.close()) + + assert artifacts[0].uri == str(tmp_path / "tail.mp4") + assert writer_calls == [ + { + "shape": (1, 2, 2, 3), + "path": tmp_path / "tail.mp4", + "fps": 8, + "layout": "thwc", + } + ] + + def test_output_sink_is_built_from_demo_spec(tmp_path: Path) -> None: def fake_writer(*args: Any, **kwargs: Any) -> Path: del args, kwargs @@ -300,6 +354,74 @@ def test_benchmark_stats_output_sink_writes_runtime_metric_samples( ] +def test_benchmark_stats_output_sink_exposes_narrow_frame_tail(tmp_path: Path) -> None: + sink = BenchmarkStatsOutputSink(output_path=tmp_path / "tail-stats.json") + sink.open(SessionInfo()) + + sink.handle_output( + 0.0, + StepResult(step_index=0, frame_count=2, metrics={"model_step_s": 0.1}), + ) + artifacts = tuple(sink.close()) + + assert artifacts[0].uri == str(tmp_path / "tail-stats.json") + payload = json.loads((tmp_path / "tail-stats.json").read_text(encoding="utf-8")) + assert payload["steps"] == [ + { + "frame_count": 2, + "metadata": {}, + "metrics": {"model_step_s": 0.1}, + "sample_count": 1, + "step_index": 0, + } + ] + + +def test_comparison_output_sink_accepts_matching_results() -> None: + result = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2)), + layout="bvtchw", + output_window=TimeWindow(start_s=0.0, end_s=1.0), + metrics={"model_step_s": 0.1}, + ) + sink = ComparisonOutputSink((result,)) + sink.open(SessionInfo(output_layout="bvtchw")) + + decision = sink.write(result) + artifacts = tuple(sink.close()) + + assert decision == OutputDecision() + assert artifacts == () + + +def test_comparison_output_sink_detects_output_regression() -> None: + expected = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.zeros((1, 1, 1, 3, 2, 2)), + layout="bvtchw", + ) + actual = StepResult.from_video_chunk( + step_index=0, + video_chunk=torch.ones((1, 1, 1, 3, 2, 2)), + layout="bvtchw", + ) + sink = ComparisonOutputSink((expected,)) + sink.open(SessionInfo(output_layout="bvtchw")) + + with pytest.raises(ComparisonOutputMismatchError, match="output tensor values"): + sink.write(actual) + + +def test_comparison_output_sink_detects_missing_result_on_close() -> None: + expected = StepResult(step_index=0, frame_count=1) + sink = ComparisonOutputSink((expected,), compare_output=False) + sink.open(SessionInfo()) + + with pytest.raises(ComparisonOutputMismatchError, match="Missing 1 expected"): + sink.close() + + def test_benchmark_output_sink_supports_stats_only(tmp_path: Path) -> None: sink = build_benchmark_output_sink(None, stats_path=tmp_path / "stats.json") sink.open(SessionInfo()) From 865f2fa14e98c097796c6d2dd61b8e856ba4370a Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 09:18:45 +0000 Subject: [PATCH 15/33] Preserve runner BaseException failures during cleanup --- flashdreams/flashdreams/demo/runner.py | 14 ++--- .../tests/test_demo_application_api.py | 52 +++++++++++++++++++ 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 67468decf..eb6884bab 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -91,7 +91,7 @@ async def run_async(self) -> RunResult: spec = self._create_spec(run_mode) context = self._create_context(host) result: RunResult | None = None - primary_error: Exception | None = None + primary_error: BaseException | None = None try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() @@ -105,7 +105,7 @@ async def run_async(self) -> RunResult: pipeline=self.pipeline or StepPipeline(), ) return result - except Exception as exc: + except BaseException as exc: primary_error = exc raise finally: @@ -126,7 +126,7 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: spec = self._create_spec(selected_run_mode) context = self._create_context(host) result: RunResult | None = None - primary_error: Exception | None = None + primary_error: BaseException | None = None try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() @@ -140,7 +140,7 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: pipeline=self.pipeline or StepPipeline(), ) return result - except Exception as exc: + except BaseException as exc: primary_error = exc raise finally: @@ -354,7 +354,7 @@ def _close_runner_resources( app: Application, owns_host: bool, run_result: RunResult | None, - primary_error: Exception | None, + primary_error: BaseException | None, ) -> None: errors: list[Exception] = [] _record_cleanup_error(errors, context.close) @@ -377,7 +377,7 @@ async def _close_runner_resources_async( app: Application, owns_host: bool, run_result: RunResult | None, - primary_error: Exception | None, + primary_error: BaseException | None, ) -> None: errors: list[Exception] = [] try: @@ -488,7 +488,7 @@ def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: def _record_cleanup_notes( - primary: Exception | None, + primary: BaseException | None, errors: Sequence[Exception], ) -> None: if primary is None: diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 06b048b0a..eaaed79f9 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -391,6 +391,27 @@ def test_runner_preserves_primary_error_when_cleanup_fails() -> None: assert app.closed +def test_runner_preserves_base_exception_when_cleanup_fails() -> None: + app = _RunnerFakeApplication( + total_steps=1, + init_error=_RunnerPrimaryBaseException("fake base failed"), + ) + + with pytest.raises( + _RunnerPrimaryBaseException, match="fake base failed" + ) as exc_info: + Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + model_id="fake-runner", + ).run() + + notes = getattr(exc_info.value, "__notes__", ()) + assert any("run metrics close failed" in note for note in notes) + assert app.closed + + def test_runner_preserves_failed_result_when_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1, fail_step=0) @@ -507,6 +528,29 @@ async def test_runner_run_async_preserves_primary_error_when_cleanup_fails() -> assert app.closed +@pytest.mark.asyncio +async def test_runner_run_async_preserves_base_exception_when_cleanup_fails() -> None: + app = _RunnerFakeApplication( + total_steps=1, + init_error=_RunnerPrimaryBaseException("fake base failed"), + ) + + with pytest.raises( + _RunnerPrimaryBaseException, match="fake base failed" + ) as exc_info: + await Runner( + io_handler=_RecordingIOHandler(), + app=app, + metrics=_FailingCloseMetricsRecorder(), + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + + notes = getattr(exc_info.value, "__notes__", ()) + assert any("run metrics close failed" in note for note in notes) + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_preserves_failed_result_when_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1, fail_step=0) @@ -1108,10 +1152,12 @@ def __init__( total_steps: int, fail_init: bool = False, fail_step: int | None = None, + init_error: BaseException | None = None, ) -> None: self.total_steps = total_steps self.fail_init = fail_init self.fail_step = fail_step + self.init_error = init_error self.launch_args: tuple[str, ...] = () self.init_thread_id: int | None = None self.session: _RunnerFakeSession | None = None @@ -1120,6 +1166,8 @@ def __init__( def init(self, launch_args: Sequence[str]) -> None: self.launch_args = tuple(launch_args) self.init_thread_id = threading.get_ident() + if self.init_error is not None: + raise self.init_error if self.fail_init: raise RuntimeError("fake init failed") @@ -1134,6 +1182,10 @@ def close(self) -> None: self.closed = True +class _RunnerPrimaryBaseException(BaseException): + pass + + class _BlockingCloseRunnerFakeApplication(_RunnerFakeApplication): def __init__( self, From d66a8b52f38bf3b8c0ef33f7f00fe06b477b422d Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 09:39:48 +0000 Subject: [PATCH 16/33] Simplify demo application discovery --- apps/t2v_demo/app.py | 25 +- apps/t2v_demo/pyproject.toml | 3 + apps/t2v_demo/tests/test_runner.py | 29 ++ flashdreams/flashdreams/demo/README.md | 317 ++++++++++++++++++ flashdreams/flashdreams/demo/__init__.py | 7 +- flashdreams/flashdreams/demo/app.py | 55 ++- flashdreams/flashdreams/plugins/__init__.py | 10 +- flashdreams/flashdreams/plugins/registry.py | 208 +++++++----- .../tests/test_demo_application_api.py | 44 ++- flashdreams/tests/test_recipe_plugins.py | 38 ++- integrations/lingbot/lingbot/demo/app.py | 76 ++--- .../omnidreams/omnidreams/demo/app.py | 62 ++-- 12 files changed, 688 insertions(+), 186 deletions(-) create mode 100644 flashdreams/flashdreams/demo/README.md diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index bee09cb67..788c6438d 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -15,6 +15,7 @@ from aiohttp import web +from flashdreams.demo import Application, DemoAdapterApplication from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, @@ -195,6 +196,28 @@ def launch_t2v( ) +def create_app(config: "T2VDemoRunnerConfig | None" = None) -> Application: + """Create the default public T2V application without CLI parsing.""" + from .runner import RUNNER_T2V + + config = RUNNER_T2V if config is None else config + adapter = make_adapter(config.backend) + scenario = _scenario(config, {}) + return DemoAdapterApplication( + adapter=adapter, + spec=_spec( + config, + adapter=adapter, + scenario=scenario, + input_mode="replay", + output=NullOutputSpec(), + ), + ) + + +createApp = create_app + + def _scenario( config: "T2VDemoRunnerConfig", overrides: dict[str, object] ) -> dict[str, object]: @@ -342,4 +365,4 @@ async def playback(_: web.Request) -> web.StreamResponse: app.router.add_get("/api/t2v/playback", playback) -__all__ = ["T2VWebRTCSessionManager", "launch_t2v"] +__all__ = ["T2VWebRTCSessionManager", "createApp", "create_app", "launch_t2v"] diff --git a/apps/t2v_demo/pyproject.toml b/apps/t2v_demo/pyproject.toml index e6b9f10e9..29698b039 100644 --- a/apps/t2v_demo/pyproject.toml +++ b/apps/t2v_demo/pyproject.toml @@ -15,6 +15,9 @@ dependencies = ["flashdreams[serving]"] [project.entry-points."flashdreams.runner_configs"] t2v = "t2v_demo.runner:RUNNER_T2V" +[project.entry-points."flashdreams.applications"] +t2v = "t2v_demo.app:create_app" + [tool.uv.sources] flashdreams = { workspace = true } diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py index 3bc17756e..afd104230 100644 --- a/apps/t2v_demo/tests/test_runner.py +++ b/apps/t2v_demo/tests/test_runner.py @@ -3,10 +3,15 @@ from __future__ import annotations +from pathlib import Path + import pytest +import tomli from t2v_demo import app from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig +from flashdreams.demo import Application + pytestmark = pytest.mark.ci_cpu @@ -15,6 +20,30 @@ def test_t2v_runner_slug_has_launch_capability() -> None: assert RUNNER_T2V.launch_capability == "t2v_demo.launch:LAUNCH_CAPABILITY" +def test_t2v_registers_application_entry_point() -> None: + pyproject_path = Path(__file__).parents[1] / "pyproject.toml" + pyproject = tomli.loads(pyproject_path.read_text()) + + assert pyproject["project"]["entry-points"]["flashdreams.applications"] == { + "t2v": "t2v_demo.app:create_app" + } + + +def test_t2v_create_app_exposes_public_application() -> None: + public_app = app.create_app( + T2VDemoRunnerConfig( + runner_name="t2v-test", + description="test", + backend="self-forcing", + prompt="A waterfall", + total_blocks=3, + ) + ) + + assert app.createApp is app.create_app + assert isinstance(public_app, Application) + + def test_runner_mp4_launch_uses_demo_entrypoint( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/flashdreams/flashdreams/demo/README.md b/flashdreams/flashdreams/demo/README.md new file mode 100644 index 000000000..a2b5b93bc --- /dev/null +++ b/flashdreams/flashdreams/demo/README.md @@ -0,0 +1,317 @@ +# FlashDreams Demo API + +`flashdreams.demo` is the public authoring layer for demos that should run +through the shared demo runtime. It gives model authors a small set of objects +to implement or compose, while the lower-level `flashdreams.runtime.demo` +package keeps ownership of the session drivers, `StepPipeline`, model worker +thread, metrics, output delivery, and cleanup. + +The goal is that a demo can support replay, MP4, null/headless, WebRTC, and +eventually native-window output without each integration reimplementing its own +loop, output handling, input plumbing, or server dispatch. + +## Main Pieces + +`Application` +: Owns app-level lifecycle. It initializes launch state, creates one model + session, and closes app resources. + +`ApplicationSession` +: Owns one model session. It reports `SessionInfo`, returns the next + `StepRequirements`, runs one `step(InferenceInput) -> StepResult`, and closes + per-session resources. The session does not write to output sinks, sleep for + backpressure, or own the loop. + +`Runner` +: Drives one session through the shared runtime. It adapts `Application` and + `IOHandler` to `run_demo_session(...)` or `run_demo_session_async(...)`, so + every mode keeps the shared worker-thread, metrics, output, cancellation, and + cleanup behavior. + +`IOHandler` +: Bundles the input, output, and stop-signal side of a run. It provides input + windows, exposes pull-style input state, emits output chunks, and reports + `should_exit()`. + +`IOHandlerServer` +: Server-shaped facade for transports such as WebRTC. WebRTC does not have a + ready IO handler until a peer connects, so a server accepts connections and + passes one `IOHandler` per session to the shared runner callback. + +`DemoAdapterApplication` +: Adapter for demos that already implement the lower-level `DemoAdapter` + contract. Most migrated demos should use this rather than implementing + `Application` from scratch. + +`create_demo_application(...)` +: Small command-app helper for demos that have parser/spec/adapter functions and + do not need a pass-through subclass. + +## Execution Model + +The runner owns the loop: + +```text +Runner + -> session.next_step_requirements() + -> IOHandler.next_window(requirements) + -> ModelInputProvider.prepare_step(...) + -> ApplicationSession.step(model_input) + -> IOHandler.emit_chunk(result) + -> metrics +``` + +This is the important boundary: the model session produces a `StepResult` and +returns it. Output delivery, backpressure, transport state, metrics, and cleanup +stay outside the model session. + +For a direct `Application` that is not backed by a `DemoAdapter`, the default +input provider passes this step payload to the session: + +```python +InferenceInput( + step={ + "step_index": requirements.step_index, + "user_window": user_window, + }, + metadata=requirements.metadata, +) +``` + +For model-specific conditioning, prefer a `DemoAdapter` with +`create_model_input_provider(...)`. That keeps IO transport details out of the +model. + +## Output Modes + +Replay-style runs use `create_replay_io_handler(...)`. + +```python +from pathlib import Path + +from flashdreams.demo import ( + DemoAdapterApplication, + FileOutputSink, + Runner, + create_replay_io_handler, +) +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec + +spec = DemoSpec( + model_id="my-demo", + input_mode="replay", + output=Mp4OutputSpec(path=Path("outputs/my-demo.mp4"), fps=30), + config=InferenceConfig(model_id="my-demo", device="cuda"), +) + +io_handler = create_replay_io_handler( + output_sink=FileOutputSink(output_path=Path("outputs/my-demo.mp4"), fps=30) +) +app = DemoAdapterApplication(adapter=MyDemoAdapter(), spec=spec) +result = Runner(io_handler=io_handler, app=app).run() +``` + +For a null/headless run, omit `output_sink` or pass a null output spec through +your command app: + +```python +io_handler = create_replay_io_handler() +result = Runner(io_handler=io_handler, app=app).run() +``` + +For metrics artifacts, add a metric tail: + +```python +from flashdreams.demo import BenchmarkStatsOutputSink + +io_handler = create_replay_io_handler( + output_sink=FileOutputSink(output_path=Path("outputs/demo.mp4"), fps=30), + metric_output_sink=BenchmarkStatsOutputSink(Path("outputs/demo-stats.json")), +) +``` + +For deterministic CI checks, use `ComparisonOutputSink` as a tail and feed it +the expected `StepResult` sequence. + +## WebRTC Shape + +WebRTC is intentionally server-shaped. A factory returns an `IOHandlerServer`, +not a ready `IOHandler`, because the transport edges only exist after a peer +connects. + +During migration, demos that still own production WebRTC serving can adapt that +serve function with `CallbackIOHandlerServer`: + +```python +from flashdreams.demo import CallbackIOHandlerServer + + +def webrtc_io_handler(args, *, context): + def serve(): + return serve_my_webrtc_demo( + spec=build_webrtc_spec(args, device=str(context.device)), + world_rank=context.world_rank, + ) + + return CallbackIOHandlerServer(serve) +``` + +Once the transport has a real shared IO handler, the server calls the same +runner callback used by replay: + +```python +server.serve(lambda handler: Runner(io_handler=handler, app=app).run()) +``` + +Do not add per-demo WebRTC managers, offer handlers, generation workers, or +WebRTC-specific runtime wrappers for new demos. Put model-specific behavior in +the adapter, provider, runtime, session, or WebRTC UI resources. + +## Pull-Based Input + +`IOHandler.get_user_input_state(modality, name)` exposes named input state as a +view over the current `UserInputWindow`. This keeps replay deterministic: a +state query and the window used for the same step always agree. + +Built-in names are: + +```python +from flashdreams.demo import InputName + +InputName.KEYBOARD +InputName.MOUSE_POSITION +InputName.MOUSE_BUTTON +InputName.HEAD_POSITION +InputName.HAND_POSITION +``` + +Keyboard state returns `KeyboardInputState`: + +```python +state = io_handler.get_user_input_state("keyboard", InputName.KEYBOARD) +if state is not None and state.is_pressed("w"): + ... +``` + +Legacy key probes such as `"key_w"` are also supported: + +```python +is_forward = io_handler.get_user_input_state("keyboard", "key_w") +``` + +For model input construction, prefer consuming the `UserInputWindow` in a +`ModelInputProvider`. The pull API is most useful for author-facing IO handlers, +tests, and simple interactive state checks. + +## Command App Helper + +For an integration CLI, define parser/spec/adapter functions and bind them with +`create_demo_application(...)`. + +```python +import argparse +from pathlib import Path + +from flashdreams.demo import create_demo_application +from flashdreams.runtime import InferenceConfig +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, NullOutputSpec + + +def parse_args(argv=None): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + + replay = subparsers.add_parser("replay") + replay.add_argument("--device", default="cuda") + replay.add_argument("--output-mode", choices=("mp4", "null"), default="mp4") + replay.add_argument("--output", type=Path) + + webrtc = subparsers.add_parser("webrtc") + webrtc.add_argument("--device", default="cuda:0") + webrtc.add_argument("--host", default="0.0.0.0") + webrtc.add_argument("--port", type=int, default=8080) + + return parser.parse_args(argv) + + +def replay_spec(args): + output = ( + NullOutputSpec() + if args.output_mode == "null" + else Mp4OutputSpec(path=args.output, fps=30) + ) + return DemoSpec( + model_id="my-demo", + input_mode="replay", + output=output, + config=InferenceConfig(model_id="my-demo", device=args.device), + ) + + +APPLICATION = create_demo_application( + parse_args=parse_args, + replay_spec=replay_spec, + replay_adapter=MyDemoAdapter, + webrtc_io_handler=webrtc_io_handler, +) + + +def main(argv=None): + APPLICATION.main(argv) +``` + +This keeps the only command branch at IO factory selection. Replay and WebRTC +both end up constructing the same public `Runner`. + +## Application Discovery + +Packages can expose public demo applications through the +`flashdreams.applications` entry-point group: + +```toml +[project.entry-points."flashdreams.applications"] +my-demo = "my_demo.app:create_app" +``` + +The referenced function should return an object satisfying `Application`: + +```python +from flashdreams.demo import Application, DemoAdapterApplication + + +def create_app() -> Application: + return DemoAdapterApplication(adapter=MyDemoAdapter(), spec=default_spec()) +``` + +Use `flashdreams.plugins.discover_applications()` to load installed demo +applications. + +## Adding A Demo + +1. Implement or adapt your model runtime. + - Existing runtime integrations should usually implement `DemoAdapter`. + - Direct demos can implement `Application` and `ApplicationSession`. +2. Keep one model step as `step(InferenceInput) -> StepResult`. + - Do not pass sinks into the session. + - Do not call output sinks from the session. + - Do not sleep for consumer backpressure inside the session. +3. Build a `DemoSpec` for each command mode. + - Put model and scenario settings in the spec. + - Keep output settings in `Mp4OutputSpec`, `NullOutputSpec`, or + `WebRTCOutputSpec`. +4. Use `create_replay_io_handler(...)` for MP4/null/replay. +5. Return an `IOHandlerServer` for WebRTC. +6. Run everything through `Runner`. +7. Add fake-model CPU tests before GPU tests. + +## Current Migration Notes + +- `create_native_window_io_handler(...)` is public, but native-window runtime + wiring is still a later migration step. +- `CallbackIOHandlerServer` exists to keep migrated demos working while + production WebRTC transports finish moving behind shared IO handlers. +- `DemoAdapterApplication` is the preferred bridge for existing runtime demo + adapters. It lets new public APIs land without rewriting every model runtime. +- Lower-level modules under `flashdreams.runtime.demo` remain the source of + truth for drivers, run modes, output sinks, metrics, and session assembly. diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 248a4dd68..6e2005b0a 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -3,7 +3,11 @@ """Public demo application authoring API.""" -from flashdreams.demo.app import DemoApplication, run_replay_application +from flashdreams.demo.app import ( + DemoApplication, + create_demo_application, + run_replay_application, +) from flashdreams.demo.application import ( Application, ApplicationSession, @@ -74,6 +78,7 @@ "WebRTCIOHandlerServer", "create_native_window_io_handler", "create_default_input_state_decoder_registry", + "create_demo_application", "create_replay_io_handler", "create_webrtc_io_handler", "input_state_from_window", diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index 4ba3bb134..9b1b9bd2e 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -7,7 +7,7 @@ import argparse import sys -from abc import ABC, abstractmethod +from collections.abc import Callable from typing import Any import torch @@ -27,9 +27,22 @@ from .runner import Runner -class DemoApplication(ABC): +class DemoApplication: """Base command application shared by model replay and WebRTC demos.""" + def __init__( + self, + *, + parse_args: Callable[[list[str] | None], argparse.Namespace] | None = None, + replay_spec: Callable[[argparse.Namespace], DemoSpec] | None = None, + replay_adapter: Callable[[], DemoAdapter] | None = None, + webrtc_io_handler: Callable[..., IOHandlerServer] | None = None, + ) -> None: + self._parse_args_fn = parse_args + self._replay_spec_fn = replay_spec + self._replay_adapter_fn = replay_adapter + self._webrtc_io_handler_fn = webrtc_io_handler + def main(self, argv: list[str] | None = None) -> None: """Parse arguments, select an IO bundle, and run the selected mode.""" configure_logging() @@ -42,17 +55,28 @@ def main(self, argv: list[str] | None = None) -> None: result = self._run_handler(args, selection) _raise_for_failed_result(result) - @abstractmethod def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: """Parse this model's command-line arguments.""" + parse_args = getattr(self, "_parse_args_fn", None) + if parse_args is None: + raise NotImplementedError("DemoApplication.parse_args is not configured.") + return parse_args(argv) - @abstractmethod def replay_spec(self, args: argparse.Namespace) -> DemoSpec: """Build the model-specific replay specification.""" + replay_spec = getattr(self, "_replay_spec_fn", None) + if replay_spec is None: + raise NotImplementedError("DemoApplication.replay_spec is not configured.") + return replay_spec(args) - @abstractmethod def replay_adapter(self) -> DemoAdapter: """Create the model-specific replay adapter.""" + replay_adapter = getattr(self, "_replay_adapter_fn", None) + if replay_adapter is None: + raise NotImplementedError( + "DemoApplication.replay_adapter is not configured." + ) + return replay_adapter() def application(self, args: argparse.Namespace) -> Application: """Create the application object consumed by the public runner.""" @@ -88,6 +112,9 @@ def webrtc_io_handler( context: Any, ) -> IOHandlerServer: """Create the WebRTC server-shaped IO factory for this model.""" + webrtc_io_handler = getattr(self, "_webrtc_io_handler_fn", None) + if webrtc_io_handler is not None: + return webrtc_io_handler(args, context=context) del args, context raise ValueError("This demo application does not support WebRTC.") @@ -106,6 +133,22 @@ def run_replay_application(*, spec: DemoSpec, adapter: DemoAdapter) -> RunResult ).run() +def create_demo_application( + *, + parse_args: Callable[[list[str] | None], argparse.Namespace], + replay_spec: Callable[[argparse.Namespace], DemoSpec], + replay_adapter: Callable[[], DemoAdapter], + webrtc_io_handler: Callable[..., IOHandlerServer] | None = None, +) -> DemoApplication: + """Create a command app from functions instead of a pass-through subclass.""" + return DemoApplication( + parse_args=parse_args, + replay_spec=replay_spec, + replay_adapter=replay_adapter, + webrtc_io_handler=webrtc_io_handler, + ) + + def _raise_for_failed_result(result: RunResult) -> None: if result.status in {"completed", "skipped"}: return @@ -116,4 +159,4 @@ def _raise_for_failed_result(result: RunResult) -> None: raise SystemExit(1) -__all__ = ["DemoApplication", "run_replay_application"] +__all__ = ["DemoApplication", "create_demo_application", "run_replay_application"] diff --git a/flashdreams/flashdreams/plugins/__init__.py b/flashdreams/flashdreams/plugins/__init__.py index 27b5cecdc..58479da1f 100644 --- a/flashdreams/flashdreams/plugins/__init__.py +++ b/flashdreams/flashdreams/plugins/__init__.py @@ -13,8 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""External-runner plugin layer (``RunnerConfig`` discovery).""" +"""Plugin discovery helpers for runners and public demo applications.""" -from flashdreams.plugins.registry import discover_runners +from flashdreams.plugins.registry import ( + discover_applications, + discover_runners, + load_plugins, +) -__all__ = ["discover_runners"] +__all__ = ["discover_applications", "discover_runners", "load_plugins"] diff --git a/flashdreams/flashdreams/plugins/registry.py b/flashdreams/flashdreams/plugins/registry.py index 779ccad51..9b45ddcc9 100644 --- a/flashdreams/flashdreams/plugins/registry.py +++ b/flashdreams/flashdreams/plugins/registry.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Discover :class:`RunnerConfig` plugins (entry-point + env-var).""" +"""Discover FlashDreams plugins registered via entry points.""" from __future__ import annotations @@ -21,14 +21,18 @@ import os import sys import traceback +from collections.abc import Callable from functools import lru_cache -from typing import cast +from typing import TYPE_CHECKING, Any, TypeVar, cast from loguru import logger from flashdreams.infra.postprocess import VideoPostProcessorConfig from flashdreams.infra.runner import RunnerConfig +if TYPE_CHECKING: + from flashdreams.demo.application import Application + if sys.version_info < (3, 10): from importlib_metadata import entry_points # type: ignore[import-not-found] else: @@ -59,6 +63,28 @@ ``cfg.runner_name``.""" +APPLICATION_ENTRY_POINT_GROUP = "flashdreams.applications" +"""Setuptools entry-point group for public demo applications.""" + +PluginT = TypeVar("PluginT") + + +def load_plugins(group: str, expected_type: type[PluginT]) -> dict[str, PluginT]: + """Load entry-point plugins and keep only values of ``expected_type``. + + Entry points may expose either a ready object or a zero-argument factory + returning one. Bad plugins are logged and skipped so a partial environment + does not break unrelated commands. + """ + plugins: dict[str, PluginT] = {} + for name, origin, value in _load_entry_point_plugins(group, expected_type): + if name in plugins: + logger.warning(f"Skipping duplicate {group} plugin {origin}.") + continue + plugins[name] = value + return plugins + + def discover_runners() -> dict[str, RunnerConfig]: """Discover externally-registered runner configs. @@ -105,58 +131,11 @@ def _accept(cfg: RunnerConfig, origin: str) -> None: runners[cfg.runner_name] = cfg origins[cfg.runner_name] = origin - # Sort entry points by name so the "first one wins" rule above is - # reproducible -- importlib.metadata gives no ordering guarantee. - discovered = sorted(entry_points(group=ENTRY_POINT_GROUP), key=lambda ep: ep.name) - for ep in discovered: - origin = f"entry point {ep.name!r} -> {ep.value}" - module_name = ep.value.split(":", 1)[0] - top_level_module = module_name.split(".", 1)[0] - try: - value = ep.load() - except ModuleNotFoundError as exc: - # Common/expected on partial installs (e.g. `uv run --project ...`): - # metadata can still expose runner entry points for integrations not - # present in the active env. If the missing module is the entry point's - # own package namespace, silently skip this plugin at debug level. - missing_name = exc.name or "" - if ( - missing_name == top_level_module - or missing_name == module_name - or missing_name.startswith(f"{top_level_module}.") - ): - logger.debug( - f"Skipping unavailable flashdreams runner {origin}: " - f"module {missing_name!r} is not installed in this environment." - ) - continue - logger.debug( - f"Failed to load flashdreams runner {origin}:\n{traceback.format_exc()}" - ) - continue - except Exception: # noqa: BLE001 - keep CLI alive on bad plugins - logger.debug( - f"Failed to load flashdreams runner {origin}:\n{traceback.format_exc()}" - ) - continue - if callable(value) and not isinstance(value, RunnerConfig): - # Allow factories that return a config (matches nerfstudio's - # env-var convention; equally useful at the entry point). - try: - value = value() - except Exception: # noqa: BLE001 - logger.warning( - f"Calling runner {origin} as a factory raised:" - f"\n{traceback.format_exc()}" - ) - continue - if not isinstance(value, RunnerConfig): - logger.warning( - f"Skipping runner {origin}: expected a RunnerConfig, " - f"got {type(value).__name__}." - ) - continue - _accept(cast(RunnerConfig, value), origin) + for _name, origin, cfg in _load_entry_point_plugins( + ENTRY_POINT_GROUP, + RunnerConfig, + ): + _accept(cfg, origin) raw = os.environ.get(ENV_VAR) if raw: @@ -188,47 +167,102 @@ def _accept(cfg: RunnerConfig, origin: str) -> None: return runners +def discover_applications() -> dict[str, "Application"]: + """Discover public demo applications registered by installed packages.""" + from flashdreams.demo.application import Application + + return load_plugins(APPLICATION_ENTRY_POINT_GROUP, Application) + + def discover_postprocess_presets() -> dict[str, VideoPostProcessorConfig]: """Discover named post-processor presets from entry points. Returns: Mapping from preset slug to :class:`VideoPostProcessorConfig`. """ - presets: dict[str, VideoPostProcessorConfig] = {} - discovered = sorted( - entry_points(group=POSTPROCESS_PRESET_GROUP), key=lambda ep: ep.name - ) + return load_plugins(POSTPROCESS_PRESET_GROUP, VideoPostProcessorConfig) + + +def _load_entry_point_plugin( + ep: Any, + *, + group: str, + origin: str, + expected_type: type[PluginT], +) -> PluginT | None: + module_name = ep.value.split(":", 1)[0] + top_level_module = module_name.split(".", 1)[0] + try: + value = ep.load() + except ModuleNotFoundError as exc: + missing_name = exc.name or "" + if ( + missing_name == top_level_module + or missing_name == module_name + or missing_name.startswith(f"{top_level_module}.") + ): + logger.debug( + f"Skipping unavailable {group} plugin {origin}: " + f"module {missing_name!r} is not installed in this environment." + ) + return None + logger.debug( + f"Failed to load {group} plugin {origin}:\n{traceback.format_exc()}" + ) + return None + except Exception: # noqa: BLE001 - keep discovery alive on bad plugins + logger.debug( + f"Failed to load {group} plugin {origin}:\n{traceback.format_exc()}" + ) + return None + if callable(value) and ( + isinstance(value, type) or not isinstance(value, expected_type) + ): + value = _call_plugin_factory(value, group=group, origin=origin) + if value is None: + return None + if not isinstance(value, expected_type): + logger.warning( + f"Skipping {group} plugin {origin}: expected a " + f"{expected_type.__name__}, got {type(value).__name__}." + ) + return None + return cast(PluginT, value) + + +def _load_entry_point_plugins( + group: str, + expected_type: type[PluginT], +) -> list[tuple[str, str, PluginT]]: + discovered = sorted(entry_points(group=group), key=lambda ep: ep.name) + plugins: list[tuple[str, str, PluginT]] = [] for ep in discovered: origin = f"entry point {ep.name!r} -> {ep.value}" - try: - value = ep.load() - except Exception: # noqa: BLE001 - keep CLI alive on bad plugins - logger.warning( - f"Failed to load postprocess preset {origin}:\n{traceback.format_exc()}" - ) - continue - if callable(value) and not isinstance(value, VideoPostProcessorConfig): - try: - value = value() - except Exception: # noqa: BLE001 - logger.warning( - f"Calling postprocess preset {origin} as a factory raised:" - f"\n{traceback.format_exc()}" - ) - continue - if not isinstance(value, VideoPostProcessorConfig): - logger.warning( - f"Skipping postprocess preset {origin}: expected a " - f"VideoPostProcessorConfig, got {type(value).__name__}." - ) - continue - if ep.name in presets: - logger.warning( - f"Skipping duplicate postprocess preset {ep.name!r} from {origin}." - ) - continue - presets[ep.name] = value - return presets + value = _load_entry_point_plugin( + ep, + group=group, + origin=origin, + expected_type=expected_type, + ) + if value is not None: + plugins.append((ep.name, origin, value)) + return plugins + + +def _call_plugin_factory( + factory: Callable[[], object], + *, + group: str, + origin: str, +) -> object | None: + try: + return factory() + except Exception: # noqa: BLE001 + logger.warning( + f"Calling {group} plugin {origin} as a factory raised:" + f"\n{traceback.format_exc()}" + ) + return None @lru_cache(maxsize=None) diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index eaaed79f9..7674d0871 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -33,6 +33,7 @@ Runner, RuntimeOutputSinkFrameAdapter, WebRTCIOHandlerServer, + create_demo_application, create_native_window_io_handler, create_replay_io_handler, create_webrtc_io_handler, @@ -91,6 +92,7 @@ def test_public_demo_contracts_are_importable() -> None: assert IOHandler.__name__ == "IOHandler" assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" + assert create_demo_application.__name__ == "create_demo_application" assert create_replay_io_handler.__name__ == "create_replay_io_handler" assert create_native_window_io_handler.__name__ == ( "create_native_window_io_handler" @@ -740,6 +742,22 @@ def test_demo_application_replay_selects_factory_and_runner() -> None: assert runtime.closed +def test_demo_application_can_be_built_from_callbacks() -> None: + adapter = _FakeDemoAdapter() + app = create_demo_application( + parse_args=_parse_replay_command, + replay_spec=_fake_replay_spec, + replay_adapter=lambda: adapter, + ) + + app.main(["replay"]) + + assert adapter.runtimes + runtime = adapter.runtimes[0] + assert runtime.session.closed + assert runtime.closed + + def test_demo_application_server_selection_does_not_build_replay_app() -> None: app = _ServerDemoApplication() @@ -757,22 +775,30 @@ def _key_event(event_type: str, key: str, timestamp_s: float) -> UserInputEvent: ) +def _parse_replay_command(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("command", choices=("replay",)) + return parser.parse_args(argv) + + +def _fake_replay_spec(args: argparse.Namespace) -> DemoSpec: + assert args.command == "replay" + return DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ) + + class _ReplayDemoApplication(DemoApplication): def __init__(self) -> None: self.adapter = _FakeDemoAdapter() def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("command", choices=("replay",)) - return parser.parse_args(argv) + return _parse_replay_command(argv) def replay_spec(self, args: argparse.Namespace) -> DemoSpec: - assert args.command == "replay" - return DemoSpec( - model_id="fake-demo", - input_mode="replay", - output=NullOutputSpec(), - ) + return _fake_replay_spec(args) def replay_adapter(self) -> "_FakeDemoAdapter": return self.adapter diff --git a/flashdreams/tests/test_recipe_plugins.py b/flashdreams/tests/test_recipe_plugins.py index 213351b51..70861b14a 100644 --- a/flashdreams/tests/test_recipe_plugins.py +++ b/flashdreams/tests/test_recipe_plugins.py @@ -37,8 +37,8 @@ ) from flashdreams.infra.config import derive_config from flashdreams.infra.runner import RunnerConfig -from flashdreams.plugins import discover_runners -from flashdreams.plugins.registry import ENV_VAR +from flashdreams.plugins import discover_applications, discover_runners +from flashdreams.plugins.registry import APPLICATION_ENTRY_POINT_GROUP, ENV_VAR from flashdreams.recipes.template.config import TEMPLATE_OFFLINE_RUNNER pytestmark = pytest.mark.ci_cpu @@ -227,6 +227,40 @@ def _capture_debug(message: str) -> None: assert any("broken-plugin" in msg for msg in debug_messages) +def test_discover_applications_uses_shared_entry_point_loader( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _PluginApplication: + def init(self, launch_args: list[str]) -> None: + del launch_args + + def create_session(self) -> object: + return object() + + def close(self) -> None: + pass + + class _FakeEntryPoint: + name = "test-application" + value = "flashdreams._test_application:create_app" + + @staticmethod + def load() -> object: + return _PluginApplication + + monkeypatch.setattr( + "flashdreams.plugins.registry.entry_points", + lambda *, group: [_FakeEntryPoint()] + if group == APPLICATION_ENTRY_POINT_GROUP + else [], + ) + + applications = discover_applications() + + assert list(applications) == ["test-application"] + assert isinstance(applications["test-application"], _PluginApplication) + + def test_discover_runners_skips_runner_name_collision( monkeypatch: pytest.MonkeyPatch, fake_plugin_module: str ) -> None: diff --git a/integrations/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py index d041a8a98..7841e773c 100644 --- a/integrations/lingbot/lingbot/demo/app.py +++ b/integrations/lingbot/lingbot/demo/app.py @@ -10,7 +10,7 @@ from typing import Any, Literal, cast from flashdreams.demo import CallbackIOHandlerServer, IOHandlerServer -from flashdreams.demo.app import DemoApplication, run_replay_application +from flashdreams.demo.app import create_demo_application, run_replay_application from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( @@ -127,49 +127,28 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return args -class LingbotDemoApplication(DemoApplication): - """Lingbot replay and WebRTC demo application.""" - - def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: - return parse_args(argv) - - def replay_spec(self, args: argparse.Namespace) -> DemoSpec: - return _replay_spec(args) - - def replay_adapter(self) -> LingbotDemoAdapter: - return LingbotDemoAdapter() - - def webrtc_io_handler( - self, - args: argparse.Namespace, - *, - context: Any, - ) -> IOHandlerServer: - def serve() -> object: - ensure_example_data_downloaded( - is_rank_zero=(context.world_rank == 0), - example_idx=args.example_idx, - ) - from .webrtc import serve_lingbot_webrtc_demo - - return serve_lingbot_webrtc_demo( - spec=_webrtc_spec( - args, - device=str(context.device), - context_parallel_size=context.world_size, - ), - world_rank=context.world_rank, - ) - - return CallbackIOHandlerServer(serve) - - -_APPLICATION = LingbotDemoApplication() +def _webrtc_io_handler( + args: argparse.Namespace, + *, + context: Any, +) -> IOHandlerServer: + def serve() -> object: + ensure_example_data_downloaded( + is_rank_zero=(context.world_rank == 0), + example_idx=args.example_idx, + ) + from .webrtc import serve_lingbot_webrtc_demo + return serve_lingbot_webrtc_demo( + spec=_webrtc_spec( + args, + device=str(context.device), + context_parallel_size=context.world_size, + ), + world_rank=context.world_rank, + ) -def main(argv: list[str] | None = None) -> None: - """Run the Lingbot demo application.""" - _APPLICATION.main(argv) + return CallbackIOHandlerServer(serve) def launch_from_runner( @@ -392,3 +371,16 @@ def _webrtc_spec( }, ), ) + + +_APPLICATION = create_demo_application( + parse_args=parse_args, + replay_spec=_replay_spec, + replay_adapter=LingbotDemoAdapter, + webrtc_io_handler=_webrtc_io_handler, +) + + +def main(argv: list[str] | None = None) -> None: + """Run the Lingbot demo application.""" + _APPLICATION.main(argv) diff --git a/integrations/omnidreams/omnidreams/demo/app.py b/integrations/omnidreams/omnidreams/demo/app.py index 084e21996..534be6b33 100644 --- a/integrations/omnidreams/omnidreams/demo/app.py +++ b/integrations/omnidreams/omnidreams/demo/app.py @@ -14,7 +14,7 @@ from omnidreams.runner import DEFAULT_EXAMPLE_DATA_UUID_1V from flashdreams.demo import CallbackIOHandlerServer, IOHandlerServer -from flashdreams.demo.app import DemoApplication, run_replay_application +from flashdreams.demo.app import create_demo_application, run_replay_application from flashdreams.infra.runner import RunnerConfig from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( @@ -125,41 +125,20 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: return args -class OmnidreamsDemoApplication(DemoApplication): - """OmniDreams replay and WebRTC demo application.""" - - def parse_args(self, argv: list[str] | None = None) -> argparse.Namespace: - return parse_args(argv) - - def replay_spec(self, args: argparse.Namespace) -> DemoSpec: - return _replay_spec(args) - - def replay_adapter(self) -> OmnidreamsDemoAdapter: - return OmnidreamsDemoAdapter() - - def webrtc_io_handler( - self, - args: argparse.Namespace, - *, - context: Any, - ) -> IOHandlerServer: - def serve() -> object: - from .webrtc import serve_omnidreams_webrtc_demo - - return serve_omnidreams_webrtc_demo( - spec=_webrtc_spec(args, device=str(context.device)), - world_rank=context.world_rank, - ) - - return CallbackIOHandlerServer(serve) - +def _webrtc_io_handler( + args: argparse.Namespace, + *, + context: Any, +) -> IOHandlerServer: + def serve() -> object: + from .webrtc import serve_omnidreams_webrtc_demo + + return serve_omnidreams_webrtc_demo( + spec=_webrtc_spec(args, device=str(context.device)), + world_rank=context.world_rank, + ) -_APPLICATION = OmnidreamsDemoApplication() - - -def main(argv: list[str] | None = None) -> None: - """Run the OmniDreams demo application.""" - _APPLICATION.main(argv) + return CallbackIOHandlerServer(serve) def launch_from_runner( @@ -397,3 +376,16 @@ def _split_paths(value: str) -> tuple[Path, ...]: def _split_strings(value: str) -> tuple[str, ...]: return tuple(part for part in value.split(",") if part) + + +_APPLICATION = create_demo_application( + parse_args=parse_args, + replay_spec=_replay_spec, + replay_adapter=OmnidreamsDemoAdapter, + webrtc_io_handler=_webrtc_io_handler, +) + + +def main(argv: list[str] | None = None) -> None: + """Run the OmniDreams demo application.""" + _APPLICATION.main(argv) From 6247dd763a8efb1a63c7bfb5c02eeb72e61195f6 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 09:50:04 +0000 Subject: [PATCH 17/33] Propagate replay schemas through demo runner Bind adapter-prepared replay inputs and source schemas onto ReplayIOHandler before shared runtime validation. This lets DemoAdapterApplication replay commands advertise provider-required raw user input schemas, including the OmniDreams Ludus keyboard event schema. Add CPU coverage for replay IO schema propagation. --- flashdreams/flashdreams/demo/io.py | 12 +++++ flashdreams/flashdreams/demo/runner.py | 17 ++++++- .../tests/test_demo_application_api.py | 47 +++++++++++++++++-- 3 files changed, 72 insertions(+), 4 deletions(-) diff --git a/flashdreams/flashdreams/demo/io.py b/flashdreams/flashdreams/demo/io.py index 9f856f7f5..c07352c2d 100644 --- a/flashdreams/flashdreams/demo/io.py +++ b/flashdreams/flashdreams/demo/io.py @@ -88,6 +88,18 @@ def __post_init__(self) -> None: self._input_source = _ReplayIOInputSource(self.replay_log or UserInputs()) self._output_sink = self.output_sink or NullOutputSink() + def configure_replay_inputs( + self, + *, + replay_log: UserInputs, + user_input_schema: UserInputSchema, + ) -> None: + """Bind adapter-prepared replay inputs before runtime validation.""" + self.user_input_schema = user_input_schema + if self.replay_log is None: + self.replay_log = replay_log + self._input_source = _ReplayIOInputSource(replay_log) + @property def run_mode(self) -> "IOHandlerRunMode": """Return the runtime run mode backing this IO handler.""" diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index eb6884bab..90091f635 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -53,7 +53,7 @@ DemoAdapterApplication, IOHandler, ) -from .io import IOHandlerRunMode +from .io import IOHandlerRunMode, ReplayIOHandler @dataclass(slots=True) @@ -95,6 +95,8 @@ async def run_async(self) -> RunResult: try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() + if isinstance(self.app, DemoAdapterApplication): + _configure_replay_io_handler(self.io_handler, scenario) adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) result = await run_demo_session_async( context=context, @@ -130,6 +132,8 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: try: self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() + if isinstance(self.app, DemoAdapterApplication): + _configure_replay_io_handler(self.io_handler, scenario) adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) result = run_demo_session( context=context, @@ -347,6 +351,17 @@ def _runner_scenario() -> PreparedScenario: return PreparedScenario(initial_inputs=InferenceInput()) +def _configure_replay_io_handler( + io_handler: IOHandler, + scenario: PreparedScenario, +) -> None: + if isinstance(io_handler, ReplayIOHandler): + io_handler.configure_replay_inputs( + replay_log=scenario.user_inputs, + user_input_schema=scenario.source_schema, + ) + + def _close_runner_resources( *, context: RunContext, diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 7674d0871..54fa2d44b 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -52,6 +52,7 @@ StepRequest, StepResult, TimeWindow, + UserInputCapability, UserInputEvent, UserInputs, UserInputSchema, @@ -258,6 +259,39 @@ def test_provider_can_pull_keyboard_state_through_replay_io_handler() -> None: ] +def test_runner_uses_prepared_source_schema_for_replay_io_validation() -> None: + schema = UserInputSchema( + capabilities=( + UserInputCapability( + event_type="keydown", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + ), + UserInputCapability( + event_type="keyup", + input_modality="keyboard", + payload_fields=frozenset({"key"}), + ), + ) + ) + io_handler = create_replay_io_handler() + adapter = _FakeDemoAdapter(source_schema=schema) + app = DemoAdapterApplication( + adapter=adapter, + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + config=InferenceConfig(model_id="fake-demo"), + ), + ) + + result = Runner(io_handler=io_handler, app=app).run() + + assert result.status == "completed" + assert io_handler.user_input_schema == schema + + def test_runtime_output_sink_frame_adapter_satisfies_frame_output_sink() -> None: output = NullOutputSink(store_results=True) output.open(SessionInfo()) @@ -1035,9 +1069,16 @@ class _FakeDemoAdapter: ) canonical_input_schema = CanonicalInputSchema() - def __init__(self) -> None: + def __init__( + self, + *, + source_schema: UserInputSchema | None = None, + user_inputs: UserInputs | None = None, + ) -> None: self.runtime = _FakeRuntime() self.runtimes: list[_FakeRuntime] = [] + self.source_schema = source_schema or UserInputSchema() + self.user_inputs = user_inputs or UserInputs() def supported_input_modes(self) -> tuple[str, ...]: return ("replay",) @@ -1061,8 +1102,8 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: assert spec.model_id == self.model_id return PreparedScenario( initial_inputs=InferenceInput(global_conditioning={"prompt": "demo"}), - user_inputs=UserInputs(), - source_schema=UserInputSchema(), + user_inputs=self.user_inputs, + source_schema=self.source_schema, ) def create_model_input_provider( From c6c2dd4e81f14c474b13df4e85e26e3966856b29 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 10:01:08 +0000 Subject: [PATCH 18/33] Migrate T2V replay to public demo runner Add a small T2VApplication/defaults facade and route T2V mp4/null launches through the public Runner and replay IO handler. This moves the batch T2V demo onto the application API while preserving the existing WebRTC serving path. Update focused T2V tests to assert public runner wiring. --- apps/t2v_demo/app.py | 200 +++++++++++++++++++++-------- apps/t2v_demo/tests/test_runner.py | 38 ++++-- 2 files changed, 176 insertions(+), 62 deletions(-) diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index 788c6438d..05dcc9d57 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -8,14 +8,22 @@ import io import json import zipfile -from dataclasses import dataclass, replace +from collections.abc import Sequence +from dataclasses import dataclass, field, replace from importlib.resources import files from pathlib import Path from typing import TYPE_CHECKING, Any, Literal from aiohttp import web -from flashdreams.demo import Application, DemoAdapterApplication +from flashdreams.demo import ( + Application, + ApplicationSession, + FileOutputSink, + InferenceSessionApplicationAdapter, + Runner, + create_replay_io_handler, +) from flashdreams.runtime import InferenceConfig from flashdreams.runtime.demo import ( DemoSpec, @@ -29,7 +37,7 @@ initialize_cuda_distributed, ) from flashdreams.runtime.demo.host import RuntimeHost -from flashdreams.runtime.demo.replay import run_replay_demo +from flashdreams.runtime.demo.run_modes import RunResult from flashdreams.serving.webrtc.demo import serve_webrtc_demo from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig @@ -42,6 +50,7 @@ FIELD_PROMPT, FIELD_TOTAL_BLOCKS, T2VDemoAdapter, + T2VRuntime, make_adapter, ) @@ -59,6 +68,69 @@ class T2VWebRTCConfig(WebRTCRuntimeConfig): warmup_timeout_s: float +@dataclass(frozen=True, slots=True) +class T2VApplicationDefaults: + """Author-facing defaults for one prompt-only text-to-video app.""" + + backend: str = "causal-forcing" + preset_id: str | None = None + prompt: str | None = None + total_blocks: int | None = None + pixel_height: int | None = None + pixel_width: int | None = None + fps: int | None = None + device: str = "cuda" + compile: bool | None = None + + +@dataclass(slots=True) +class T2VApplication: + """Small public app facade over the selected T2V runtime preset.""" + + defaults: T2VApplicationDefaults = field(default_factory=T2VApplicationDefaults) + _runtimes: list[T2VRuntime] = field(default_factory=list, init=False, repr=False) + + def init(self, launch_args: Sequence[str]) -> None: + if launch_args: + raise ValueError( + "T2VApplication does not support launch arguments; pass defaults " + "when constructing the app." + ) + resolve_backend(self.defaults.backend).resolve_runner(self.defaults.preset_id) + + def create_session(self) -> ApplicationSession: + adapter = make_adapter(self.defaults.backend) + scenario = _scenario(self.defaults) + spec = _spec( + self.defaults, + adapter=adapter, + scenario=scenario, + input_mode="replay", + output=NullOutputSpec(), + ) + if spec.config is None: + raise RuntimeError("T2V DemoSpec.config was not initialized.") + prepared = adapter.prepare_scenario(spec) + runtime = adapter.create_runtime(spec.config) + self._runtimes.append(runtime) + return InferenceSessionApplicationAdapter( + runtime.start_session(prepared.initial_inputs) + ) + + def close(self) -> None: + errors: list[Exception] = [] + while self._runtimes: + runtime = self._runtimes.pop() + try: + runtime.close() + except Exception as exc: + errors.append(exc) + if errors: + raise RuntimeError( + f"T2VApplication.close failed for {len(errors)} runtime(s)." + ) from errors[0] + + class T2VWebRTCSessionManager(BaseWebRTCSessionManager[Any, T2VWebRTCConfig]): """Shared manager with a prompt update for the next browser session.""" @@ -94,8 +166,8 @@ def launch_t2v( configure_logging() scenario_overrides = scenario_overrides or {} output_overrides = output_overrides or {} - adapter = make_adapter(config.backend) - scenario = _scenario(config, scenario_overrides) + defaults = _defaults_from_config(config, scenario_overrides) + scenario = _scenario(defaults) if mode == "mp4" or mode == "null": output = _replay_output( mode=mode, @@ -106,22 +178,13 @@ def launch_t2v( output_overrides.get("fps", scenario[FIELD_FPS]), name="fps" ), ) - result = run_replay_demo( - spec=_spec( - config, - adapter=adapter, - scenario=scenario, - input_mode="replay", - output=output, - ), - adapter=adapter, + return _run_replay_application( + app=T2VApplication(defaults=defaults), + output=output, ) - if result.status != "completed": - reason = result.reason or str(result.error) or "T2V replay failed." - raise RuntimeError(reason) - return result context = initialize_cuda_distributed(default_device=config.device) + adapter = make_adapter(defaults.backend) output = WebRTCOutputSpec( host=str(host or output_overrides.get("host", "0.0.0.0")), port=_int_value( @@ -151,12 +214,11 @@ def launch_t2v( preload_name="FlashDreams T2V", ) spec = _spec( - config, + replace(defaults, device=str(context.device)), adapter=adapter, scenario=scenario, input_mode="webrtc", output=output, - device=str(context.device), ) prepared = adapter.prepare_scenario(spec) inference_config = spec.config @@ -201,36 +263,38 @@ def create_app(config: "T2VDemoRunnerConfig | None" = None) -> Application: from .runner import RUNNER_T2V config = RUNNER_T2V if config is None else config - adapter = make_adapter(config.backend) - scenario = _scenario(config, {}) - return DemoAdapterApplication( - adapter=adapter, - spec=_spec( - config, - adapter=adapter, - scenario=scenario, - input_mode="replay", - output=NullOutputSpec(), - ), - ) + return T2VApplication(defaults=_defaults_from_config(config, {})) createApp = create_app -def _scenario( - config: "T2VDemoRunnerConfig", overrides: dict[str, object] -) -> dict[str, object]: - runner = resolve_backend(config.backend).resolve_runner(config.preset_id) +def _defaults_from_config( + config: "T2VDemoRunnerConfig", + overrides: dict[str, object], +) -> T2VApplicationDefaults: + def value(name: str) -> object: + return getattr(config, name) if overrides.get(name) is None else overrides[name] + + return T2VApplicationDefaults( + backend=config.backend, + preset_id=config.preset_id, + prompt=_optional_str(value(FIELD_PROMPT)), + total_blocks=_optional_int(value(FIELD_TOTAL_BLOCKS)), + pixel_height=_optional_int(value(FIELD_PIXEL_HEIGHT)), + pixel_width=_optional_int(value(FIELD_PIXEL_WIDTH)), + fps=_optional_int(value(FIELD_FPS)), + device=config.device, + compile=config.compile, + ) + + +def _scenario(defaults: T2VApplicationDefaults) -> dict[str, object]: + runner = resolve_backend(defaults.backend).resolve_runner(defaults.preset_id) def value(name: str, default: object) -> object: - overridden = overrides.get(name) - configured = getattr(config, name) - return ( - default - if overridden is None and configured is None - else (configured if overridden is None else overridden) - ) + configured = getattr(defaults, name) + return default if configured is None else configured return { FIELD_PROMPT: value(FIELD_PROMPT, runner.prompt), @@ -242,25 +306,24 @@ def value(name: str, default: object) -> object: def _spec( - config: "T2VDemoRunnerConfig", + defaults: T2VApplicationDefaults, *, adapter: T2VDemoAdapter, scenario: dict[str, object], input_mode: Literal["replay", "webrtc"], output: Mp4OutputSpec | NullOutputSpec | WebRTCOutputSpec, - device: str | None = None, ) -> DemoSpec: return DemoSpec( model_id=adapter.model_id, - preset_id=config.preset_id or adapter.backend.default_preset_name, + preset_id=defaults.preset_id or adapter.backend.default_preset_name, input_mode=input_mode, scenario=scenario, output=output, config=InferenceConfig( model_id=adapter.model_id, - preset_id=config.preset_id or adapter.backend.default_preset_name, - device=device or config.device, - compile=config.compile, + preset_id=defaults.preset_id or adapter.backend.default_preset_name, + device=defaults.device, + compile=defaults.compile, runtime_options={"backend": adapter.backend.key}, ), ) @@ -298,6 +361,36 @@ def _float_value(value: object, *, name: str) -> float: raise TypeError(f"{name} must be convertible to float, got {type(value).__name__}.") +def _optional_int(value: object) -> int | None: + return None if value is None else _int_value(value, name="value") + + +def _optional_str(value: object) -> str | None: + return None if value is None else str(value) + + +def _run_replay_application( + *, + app: Application, + output: Mp4OutputSpec | NullOutputSpec, +) -> RunResult: + output_sink = None + if isinstance(output, Mp4OutputSpec): + output_sink = FileOutputSink( + output_path=output.path, + fps=output.fps, + output_layout=output.output_layout, + ) + result = Runner( + io_handler=create_replay_io_handler(output_sink=output_sink), + app=app, + ).run() + if result.status != "completed": + reason = result.reason or str(result.error) or "T2V replay failed." + raise RuntimeError(reason) + return result + + def _configure_app( app: web.Application, *, @@ -365,4 +458,11 @@ async def playback(_: web.Request) -> web.StreamResponse: app.router.add_get("/api/t2v/playback", playback) -__all__ = ["T2VWebRTCSessionManager", "createApp", "create_app", "launch_t2v"] +__all__ = [ + "T2VApplication", + "T2VApplicationDefaults", + "T2VWebRTCSessionManager", + "createApp", + "create_app", + "launch_t2v", +] diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py index afd104230..47e945604 100644 --- a/apps/t2v_demo/tests/test_runner.py +++ b/apps/t2v_demo/tests/test_runner.py @@ -10,7 +10,8 @@ from t2v_demo import app from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig -from flashdreams.demo import Application +from flashdreams.demo import Application, FileOutputSink, ReplayIOHandler +from flashdreams.runtime.demo import RunResult pytestmark = pytest.mark.ci_cpu @@ -42,18 +43,26 @@ def test_t2v_create_app_exposes_public_application() -> None: assert app.createApp is app.create_app assert isinstance(public_app, Application) + assert isinstance(public_app, app.T2VApplication) + assert public_app.defaults.backend == "self-forcing" + assert public_app.defaults.prompt == "A waterfall" + assert public_app.defaults.total_blocks == 3 def test_runner_mp4_launch_uses_demo_entrypoint( monkeypatch: pytest.MonkeyPatch, ) -> None: - captured = [] + captured: dict[str, object] = {} - def fake_replay_demo(*, spec: object, adapter: object) -> object: - captured.append((spec, adapter)) - return type("Result", (), {"status": "completed"})() + class FakeRunner: + def __init__(self, *, io_handler: object, app: object) -> None: + captured["io_handler"] = io_handler + captured["app"] = app - monkeypatch.setattr(app, "run_replay_demo", fake_replay_demo) + def run(self) -> RunResult: + return RunResult(status="completed") + + monkeypatch.setattr(app, "Runner", FakeRunner) config = T2VDemoRunnerConfig( runner_name="t2v", description="test", @@ -68,9 +77,14 @@ def fake_replay_demo(*, spec: object, adapter: object) -> object: output_overrides={"path": "outputs/test.mp4", "fps": 24}, ) - spec, _adapter = captured[0] - assert spec.input_mode == "replay" - assert spec.scenario["prompt"] == "A waterfall" - assert spec.scenario["total_blocks"] == 3 - assert str(spec.output.path) == "outputs/test.mp4" - assert spec.output.fps == 24 + public_app = captured["app"] + io_handler = captured["io_handler"] + assert isinstance(public_app, app.T2VApplication) + assert isinstance(io_handler, ReplayIOHandler) + assert public_app.defaults.prompt == "A waterfall" + assert public_app.defaults.total_blocks == 3 + output_sink = io_handler.output_sink + assert isinstance(output_sink, FileOutputSink) + assert str(output_sink.output_path) == "outputs/test.mp4" + assert output_sink.fps == 24 + assert output_sink.output_layout == "tchw" From a85048cbe7cb226c730d4a092ba481e657d90202 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 10:24:24 +0000 Subject: [PATCH 19/33] Extract neutral T2V demo shell --- apps/t2v/__init__.py | 42 +++ apps/t2v/pyproject.toml | 20 ++ apps/t2v/t2v.py | 526 +++++++++++++++++++++++++++++ apps/t2v/tests/test_t2v_shell.py | 95 ++++++ apps/t2v_demo/app.py | 155 ++------- apps/t2v_demo/pyproject.toml | 3 +- apps/t2v_demo/runtime.py | 331 +++--------------- apps/t2v_demo/tests/test_runner.py | 42 ++- uv.lock | 18 +- 9 files changed, 813 insertions(+), 419 deletions(-) create mode 100644 apps/t2v/__init__.py create mode 100644 apps/t2v/pyproject.toml create mode 100644 apps/t2v/t2v.py create mode 100644 apps/t2v/tests/test_t2v_shell.py diff --git a/apps/t2v/__init__.py b/apps/t2v/__init__.py new file mode 100644 index 000000000..920e1595c --- /dev/null +++ b/apps/t2v/__init__.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared text-to-video demo application shell.""" + +from t2v.t2v import ( + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + T2VDemoAdapter, + T2VInputProvider, + T2VModelConfig, + T2VRunDefaults, + T2VRuntime, + T2VScenario, + T2VSession, + create_t2v_application, + create_t2v_spec, + run_t2v_replay_application, + t2v_scenario_mapping, +) + +__all__ = [ + "FIELD_FPS", + "FIELD_PIXEL_HEIGHT", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "T2VDemoAdapter", + "T2VInputProvider", + "T2VModelConfig", + "T2VRunDefaults", + "T2VRuntime", + "T2VScenario", + "T2VSession", + "create_t2v_application", + "create_t2v_spec", + "run_t2v_replay_application", + "t2v_scenario_mapping", +] diff --git a/apps/t2v/pyproject.toml b/apps/t2v/pyproject.toml new file mode 100644 index 000000000..e48918087 --- /dev/null +++ b/apps/t2v/pyproject.toml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-t2v" +version = "0.1.0" +description = "Shared model-neutral text-to-video demo application shell" +requires-python = ">=3.10" +dependencies = ["flashdreams"] + +[tool.uv.sources] +flashdreams = { workspace = true } + +[tool.setuptools] +packages = ["t2v"] +package-dir = { t2v = "." } diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py new file mode 100644 index 000000000..585b3cbf9 --- /dev/null +++ b/apps/t2v/t2v.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Model-neutral text-to-video demo shell.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import torch + +from flashdreams.demo import ( + Application, + DemoAdapterApplication, + FileOutputSink, + Runner, + create_replay_io_handler, +) +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.video_output import VideoOutputStream +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputField, + ModelAdapter, + StepRequest, +) +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + OutputSpec, + PreparedScenario, +) +from flashdreams.runtime.demo.outputs import SessionInfo +from flashdreams.runtime.demo.run_modes import RunResult +from flashdreams.runtime.demo.session_inputs import ( + PreparedStep, + ProviderCapabilities, + UserInputWindow, +) +from flashdreams.runtime.interfaces import InferenceSession +from flashdreams.runtime.types import StepResult +from flashdreams.runtime.video_output import Mp4VideoOutputTarget + +FIELD_PROMPT = "prompt" +FIELD_TOTAL_BLOCKS = "total_blocks" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" +FIELD_FPS = "fps" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class T2VModelConfig: + """One integration-owned prompt-only text-to-video model entry.""" + + model_id: str + preset_id: str | None + pipeline: Any + prompt: str + total_blocks: int + pixel_height: int + pixel_width: int + fps: int + runtime_options: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("T2VModelConfig.model_id must be non-empty.") + if self.preset_id is not None and not self.preset_id.strip(): + raise ValueError("T2VModelConfig.preset_id must be non-empty when set.") + if not self.prompt.strip(): + raise ValueError("T2VModelConfig.prompt must be non-empty.") + _validate_positive_int(self.total_blocks, name=FIELD_TOTAL_BLOCKS) + _validate_positive_int(self.pixel_height, name=FIELD_PIXEL_HEIGHT) + _validate_positive_int(self.pixel_width, name=FIELD_PIXEL_WIDTH) + _validate_positive_int(self.fps, name=FIELD_FPS) + object.__setattr__( + self, "runtime_options", freeze_mapping(self.runtime_options) + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class T2VRunDefaults: + """Launch-time overrides shared by T2V replay and WebRTC modes.""" + + prompt: str | None = None + total_blocks: int | None = None + pixel_height: int | None = None + pixel_width: int | None = None + fps: int | None = None + device: str = "cuda" + compile: bool | None = None + runtime_options: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.device.strip(): + raise ValueError("T2VRunDefaults.device must be non-empty.") + if self.prompt is not None and not self.prompt.strip(): + raise ValueError("T2VRunDefaults.prompt must be non-empty when set.") + _validate_optional_positive_int(self.total_blocks, name=FIELD_TOTAL_BLOCKS) + _validate_optional_positive_int(self.pixel_height, name=FIELD_PIXEL_HEIGHT) + _validate_optional_positive_int(self.pixel_width, name=FIELD_PIXEL_WIDTH) + _validate_optional_positive_int(self.fps, name=FIELD_FPS) + object.__setattr__( + self, "runtime_options", freeze_mapping(self.runtime_options) + ) + + +@dataclass(frozen=True, kw_only=True, slots=True) +class T2VScenario: + """Prompt and output geometry for a finite text-to-video rollout.""" + + prompt: str + total_blocks: int + pixel_height: int + pixel_width: int + fps: int + + +class T2VDemoAdapter(ModelAdapter): + """Model adapter shared by replay and WebRTC T2V launch paths.""" + + inference_input_schema = InferenceInputSchema( + global_conditioning_fields=(InputField(name=FIELD_PROMPT),), + description="Text-to-video prompt and rollout settings.", + ) + canonical_input_schema = CanonicalInputSchema() + + def __init__(self, *, model: T2VModelConfig) -> None: + self.model = model + + @property + def model_id(self) -> str: + return self.model.model_id + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay", "webrtc") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "null", "webrtc") + + def default_input_mapping(self) -> IdentityInputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + if config.model_id != self.model.model_id: + raise ValueError( + f"Expected model_id={self.model.model_id!r}, " + f"got {config.model_id!r}." + ) + if ( + self.model.preset_id is not None + and config.preset_id != self.model.preset_id + ): + raise ValueError( + f"Expected preset_id={self.model.preset_id!r}, " + f"got {config.preset_id!r}." + ) + for name, expected in self.model.runtime_options.items(): + if config.runtime_options.get(name) != expected: + raise ValueError( + f"Expected runtime option {name}={expected!r}, " + f"got {config.runtime_options.get(name)!r}." + ) + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + scenario = _scenario_from_value(spec.scenario, self.model) + return PreparedScenario( + initial_inputs=InferenceInput( + global_conditioning={ + FIELD_PROMPT: scenario.prompt, + FIELD_TOTAL_BLOCKS: scenario.total_blocks, + FIELD_PIXEL_HEIGHT: scenario.pixel_height, + FIELD_PIXEL_WIDTH: scenario.pixel_width, + FIELD_FPS: scenario.fps, + } + ) + ) + + def create_runtime(self, config: InferenceConfig) -> "T2VRuntime": + self.validate_config(config) + return T2VRuntime(config=config, model=self.model) + + def create_model_input_provider( + self, spec: DemoSpec, scenario: PreparedScenario + ) -> "T2VInputProvider": + """Supply fixed prompt conditioning to every shared-demo step.""" + del spec + return T2VInputProvider(initial_inputs=scenario.initial_inputs) + + +class T2VInputProvider: + """No-control input provider for finite prompt-only generation.""" + + capabilities = ProviderCapabilities( + supports_realtime_clock=True, + supports_recorded_input=True, + deterministic_given_inputs=True, + ) + + def __init__(self, *, initial_inputs: InferenceInput) -> None: + self._initial_inputs = initial_inputs + + def prepare_initial_input(self) -> InferenceInput: + return self._initial_inputs + + def prepare_step( + self, *, request: Any, user_window: UserInputWindow + ) -> PreparedStep: + del request, user_window + return PreparedStep(inference_input=InferenceInput()) + + def reset(self, inputs: InferenceInput | None = None) -> None: + if inputs is not None: + self._initial_inputs = inputs + + def close(self) -> None: + pass + + +class T2VRuntime: + """One heavyweight selected pipeline, reusable across demo sessions.""" + + def __init__(self, *, config: InferenceConfig, model: T2VModelConfig) -> None: + self.config = config + self.model = model + pipeline_config = model.pipeline + if config.compile is not None: + from flashdreams.infra.config import derive_config + + pipeline_config = derive_config( + base_config=pipeline_config, + diffusion_model={"transformer": {"compile_network": config.compile}}, + ) + self.pipeline = pipeline_config.setup().to(config.device or "cuda").eval() + self._latest_artifact: tuple[Path, T2VScenario] | None = None + + def blocks_for_duration(self, duration_s: float, *, fps: int) -> int: + """Return enough autoregressive chunks to reach the requested duration.""" + target_frames = int(duration_s * fps) + frames = 0 + index = 0 + while frames < target_frames: + frames += int(self.pipeline.get_num_output_frames(index)) + index += 1 + return index + + def record_artifact(self, path: Path, scenario: T2VScenario) -> None: + self._latest_artifact = (path, scenario) + + @property + def latest_artifact(self) -> tuple[Path, T2VScenario] | None: + return self._latest_artifact + + def start_session(self, inputs: InferenceInput) -> "T2VSession": + return T2VSession( + pipeline=self.pipeline, scenario=_scenario_from_inputs(inputs), runtime=self + ) + + def close(self) -> None: + close = getattr(self.pipeline, "close", None) + if callable(close): + close() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class T2VSession(InferenceSession): + """A cache-isolated T2V session that yields chunks as they are generated.""" + + def __init__( + self, *, pipeline: Any, scenario: T2VScenario, runtime: T2VRuntime + ) -> None: + self.pipeline = pipeline + self.scenario = scenario + self._runtime = runtime + self._artifact_path = Path("outputs/t2v-webrtc") / f"{uuid4()}.mp4" + self._artifact_path.parent.mkdir(parents=True, exist_ok=True) + # Transitional support for the existing WebRTC download endpoint. + # Replay modes still deliver primary output through the shared OutputSink. + self._artifact_output = Mp4VideoOutputTarget( + output_path=self._artifact_path, fps=scenario.fps, output_layout="tchw" + ) + self._artifact_output.open() + self._step_index = 0 + self._closed = False + self._output_stream = VideoOutputStream( + postprocess_stream=None, output_layout="tchw" + ) + assert isinstance(pipeline.decoder, StreamingVideoDecoder) + ratio = pipeline.decoder.spatial_compression_ratio + if scenario.pixel_height % ratio or scenario.pixel_width % ratio: + raise ValueError( + "T2V dimensions must be divisible by the decoder spatial " + "compression ratio." + ) + self._cache = pipeline.initialize_cache( + text=[scenario.prompt], + image=None, + height=scenario.pixel_height // ratio, + width=scenario.pixel_width // ratio, + ) + + def session_info(self) -> SessionInfo: + return SessionInfo( + output_layout="tchw", + metadata={ + FIELD_PROMPT: self.scenario.prompt, + FIELD_TOTAL_BLOCKS: self.scenario.total_blocks, + FIELD_PIXEL_HEIGHT: self.scenario.pixel_height, + FIELD_PIXEL_WIDTH: self.scenario.pixel_width, + FIELD_FPS: self.scenario.fps, + }, + ) + + def next_step_request(self) -> StepRequest | None: + if self._closed or self._step_index >= self.scenario.total_blocks: + return None + return StepRequest(step_index=self._step_index) + + def step(self, inputs: InferenceInput) -> StepResult: + del inputs + if self._closed: + raise RuntimeError("T2V session is closed.") + index = self._step_index + video = self.pipeline.generate(autoregressive_index=index, cache=self._cache) + stats = self.pipeline.finalize(autoregressive_index=index, cache=self._cache) + self._step_index += 1 + result = self._output_stream.process( + video, + autoregressive_index=index, + metrics=stats, + metadata={FIELD_PROMPT: self.scenario.prompt}, + ) + self._artifact_output.write(result) + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + if inputs is not None and _scenario_from_inputs(inputs) != self.scenario: + raise ValueError( + "Create a new T2V session to change the prompt or dimensions." + ) + raise RuntimeError( + "T2V sessions are finite; create a new session instead of reset()." + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + artifacts = self._artifact_output.close() + if artifacts: + self._runtime.record_artifact(self._artifact_path, self.scenario) + + +def create_t2v_application( + *, + model: T2VModelConfig, + defaults: T2VRunDefaults | None = None, + input_mode: str = "replay", + output: OutputSpec | None = None, +) -> Application: + """Create a public application for one integration-owned T2V model.""" + output = output or NullOutputSpec() + spec = create_t2v_spec( + model=model, + defaults=defaults, + input_mode=input_mode, + output=output, + ) + return DemoAdapterApplication(adapter=T2VDemoAdapter(model=model), spec=spec) + + +def create_t2v_spec( + *, + model: T2VModelConfig, + defaults: T2VRunDefaults | None = None, + input_mode: str, + output: OutputSpec, +) -> DemoSpec: + """Build the shared runtime spec for one T2V run.""" + defaults = defaults or T2VRunDefaults() + return DemoSpec( + model_id=model.model_id, + preset_id=model.preset_id, + input_mode=input_mode, + scenario=t2v_scenario_mapping(model=model, defaults=defaults), + output=output, + config=InferenceConfig( + model_id=model.model_id, + preset_id=model.preset_id, + device=defaults.device, + compile=defaults.compile, + runtime_options={ + **model.runtime_options, + **defaults.runtime_options, + }, + ), + ) + + +def run_t2v_replay_application( + *, + model: T2VModelConfig, + defaults: T2VRunDefaults | None, + output: Mp4OutputSpec | NullOutputSpec, +) -> RunResult: + """Run finite T2V replay through the public runner and replay IO handler.""" + output_sink = None + if isinstance(output, Mp4OutputSpec): + output_sink = FileOutputSink( + output_path=Path(output.path), + fps=output.fps, + output_layout=output.output_layout, + ) + result = Runner( + io_handler=create_replay_io_handler(output_sink=output_sink), + app=create_t2v_application(model=model, defaults=defaults, output=output), + ).run() + if result.status != "completed": + reason = result.reason or str(result.error) or "T2V replay failed." + raise RuntimeError(reason) + return result + + +def t2v_scenario_mapping( + *, model: T2VModelConfig, defaults: T2VRunDefaults | None = None +) -> dict[str, object]: + """Return runtime scenario values after applying launch overrides.""" + defaults = defaults or T2VRunDefaults() + return { + FIELD_PROMPT: model.prompt if defaults.prompt is None else defaults.prompt, + FIELD_TOTAL_BLOCKS: ( + model.total_blocks + if defaults.total_blocks is None + else defaults.total_blocks + ), + FIELD_PIXEL_HEIGHT: ( + model.pixel_height + if defaults.pixel_height is None + else defaults.pixel_height + ), + FIELD_PIXEL_WIDTH: ( + model.pixel_width if defaults.pixel_width is None else defaults.pixel_width + ), + FIELD_FPS: model.fps if defaults.fps is None else defaults.fps, + } + + +def _scenario_from_value(value: Any, model: T2VModelConfig) -> T2VScenario: + source = value if isinstance(value, dict) else {} + prompt = str(source.get(FIELD_PROMPT, model.prompt)).strip() + if not prompt: + raise ValueError("A non-empty text-to-video prompt is required.") + return T2VScenario( + prompt=prompt, + total_blocks=_int_value(source.get(FIELD_TOTAL_BLOCKS, model.total_blocks)), + pixel_height=_int_value(source.get(FIELD_PIXEL_HEIGHT, model.pixel_height)), + pixel_width=_int_value(source.get(FIELD_PIXEL_WIDTH, model.pixel_width)), + fps=_int_value(source.get(FIELD_FPS, model.fps)), + ) + + +def _scenario_from_inputs(inputs: InferenceInput) -> T2VScenario: + source = inputs.global_conditioning + return T2VScenario( + prompt=str(source[FIELD_PROMPT]), + total_blocks=_int_value(source[FIELD_TOTAL_BLOCKS]), + pixel_height=_int_value(source[FIELD_PIXEL_HEIGHT]), + pixel_width=_int_value(source[FIELD_PIXEL_WIDTH]), + fps=_int_value(source[FIELD_FPS]), + ) + + +def _int_value(value: object) -> int: + if isinstance(value, bool): + raise TypeError("Expected an integer, not bool.") + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + if isinstance(value, str): + return int(value) + raise TypeError(f"Expected an integer-compatible value, got {type(value).__name__}.") + + +def _validate_positive_int(value: int, *, name: str) -> None: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer.") + if value <= 0: + raise ValueError(f"{name} must be > 0.") + + +def _validate_optional_positive_int(value: int | None, *, name: str) -> None: + if value is None: + return + _validate_positive_int(value, name=name) + + +__all__ = [ + "FIELD_FPS", + "FIELD_PIXEL_HEIGHT", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "T2VDemoAdapter", + "T2VInputProvider", + "T2VModelConfig", + "T2VRunDefaults", + "T2VRuntime", + "T2VScenario", + "T2VSession", + "create_t2v_application", + "create_t2v_spec", + "run_t2v_replay_application", + "t2v_scenario_mapping", +] diff --git a/apps/t2v/tests/test_t2v_shell.py b/apps/t2v/tests/test_t2v_shell.py new file mode 100644 index 000000000..d0bd02079 --- /dev/null +++ b/apps/t2v/tests/test_t2v_shell.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import inspect + +import pytest + +from flashdreams.demo import DemoAdapterApplication +from flashdreams.runtime.demo import Mp4OutputSpec, NullOutputSpec +from t2v.t2v import ( + T2VDemoAdapter, + T2VModelConfig, + T2VRunDefaults, + create_t2v_application, + create_t2v_spec, + t2v_scenario_mapping, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_t2v_shell_builds_spec_from_model_defaults() -> None: + model = _fake_model() + + spec = create_t2v_spec( + model=model, + defaults=T2VRunDefaults(device="cuda:0"), + input_mode="replay", + output=NullOutputSpec(), + ) + + assert spec.model_id == "fake-t2v" + assert spec.preset_id == "fake-preset" + assert spec.config is not None + assert spec.config.device == "cuda:0" + assert spec.config.runtime_options["family"] == "fake" + assert spec.scenario == { + "prompt": "A test prompt", + "total_blocks": 2, + "pixel_height": 32, + "pixel_width": 64, + "fps": 8, + } + + +def test_t2v_shell_applies_launch_overrides() -> None: + model = _fake_model() + + scenario = t2v_scenario_mapping( + model=model, + defaults=T2VRunDefaults(prompt="Override", total_blocks=4), + ) + + assert scenario["prompt"] == "Override" + assert scenario["total_blocks"] == 4 + assert scenario["pixel_height"] == 32 + + +def test_t2v_shell_creates_demo_adapter_application() -> None: + public_app = create_t2v_application( + model=_fake_model(), + defaults=T2VRunDefaults(), + output=Mp4OutputSpec(path="outputs/fake.mp4", fps=8, output_layout="tchw"), + ) + + assert isinstance(public_app, DemoAdapterApplication) + assert isinstance(public_app.adapter, T2VDemoAdapter) + assert public_app.spec.output.mode == "mp4" + + +def test_t2v_shell_has_no_legacy_backend_imports() -> None: + import t2v.t2v as t2v_shell + + source = inspect.getsource(t2v_shell) + + assert "t2v_demo" not in source + assert "causal_forcing" not in source + assert "self_forcing" not in source + assert "cosmos_predict2" not in source + + +def _fake_model() -> T2VModelConfig: + return T2VModelConfig( + model_id="fake-t2v", + preset_id="fake-preset", + pipeline=object(), + prompt="A test prompt", + total_blocks=2, + pixel_height=32, + pixel_width=64, + fps=8, + runtime_options={"family": "fake"}, + ) diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index 05dcc9d57..e29ec0a51 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -8,23 +8,21 @@ import io import json import zipfile -from collections.abc import Sequence -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from importlib.resources import files from pathlib import Path from typing import TYPE_CHECKING, Any, Literal from aiohttp import web - -from flashdreams.demo import ( - Application, - ApplicationSession, - FileOutputSink, - InferenceSessionApplicationAdapter, - Runner, - create_replay_io_handler, +from t2v.t2v import ( + T2VRunDefaults, + create_t2v_application, + create_t2v_spec, + run_t2v_replay_application, + t2v_scenario_mapping, ) -from flashdreams.runtime import InferenceConfig + +from flashdreams.demo import Application from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, @@ -37,12 +35,11 @@ initialize_cuda_distributed, ) from flashdreams.runtime.demo.host import RuntimeHost -from flashdreams.runtime.demo.run_modes import RunResult from flashdreams.serving.webrtc.demo import serve_webrtc_demo from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig -from .backends import backend_metadata, resolve_backend +from .backends import backend_metadata from .runtime import ( FIELD_FPS, FIELD_PIXEL_HEIGHT, @@ -50,8 +47,8 @@ FIELD_PROMPT, FIELD_TOTAL_BLOCKS, T2VDemoAdapter, - T2VRuntime, make_adapter, + model_from_backend, ) if TYPE_CHECKING: @@ -68,67 +65,12 @@ class T2VWebRTCConfig(WebRTCRuntimeConfig): warmup_timeout_s: float -@dataclass(frozen=True, slots=True) -class T2VApplicationDefaults: - """Author-facing defaults for one prompt-only text-to-video app.""" +@dataclass(frozen=True, kw_only=True, slots=True) +class T2VApplicationDefaults(T2VRunDefaults): + """Legacy ``t2v`` command defaults plus temporary backend selection.""" backend: str = "causal-forcing" preset_id: str | None = None - prompt: str | None = None - total_blocks: int | None = None - pixel_height: int | None = None - pixel_width: int | None = None - fps: int | None = None - device: str = "cuda" - compile: bool | None = None - - -@dataclass(slots=True) -class T2VApplication: - """Small public app facade over the selected T2V runtime preset.""" - - defaults: T2VApplicationDefaults = field(default_factory=T2VApplicationDefaults) - _runtimes: list[T2VRuntime] = field(default_factory=list, init=False, repr=False) - - def init(self, launch_args: Sequence[str]) -> None: - if launch_args: - raise ValueError( - "T2VApplication does not support launch arguments; pass defaults " - "when constructing the app." - ) - resolve_backend(self.defaults.backend).resolve_runner(self.defaults.preset_id) - - def create_session(self) -> ApplicationSession: - adapter = make_adapter(self.defaults.backend) - scenario = _scenario(self.defaults) - spec = _spec( - self.defaults, - adapter=adapter, - scenario=scenario, - input_mode="replay", - output=NullOutputSpec(), - ) - if spec.config is None: - raise RuntimeError("T2V DemoSpec.config was not initialized.") - prepared = adapter.prepare_scenario(spec) - runtime = adapter.create_runtime(spec.config) - self._runtimes.append(runtime) - return InferenceSessionApplicationAdapter( - runtime.start_session(prepared.initial_inputs) - ) - - def close(self) -> None: - errors: list[Exception] = [] - while self._runtimes: - runtime = self._runtimes.pop() - try: - runtime.close() - except Exception as exc: - errors.append(exc) - if errors: - raise RuntimeError( - f"T2VApplication.close failed for {len(errors)} runtime(s)." - ) from errors[0] class T2VWebRTCSessionManager(BaseWebRTCSessionManager[Any, T2VWebRTCConfig]): @@ -178,13 +120,14 @@ def launch_t2v( output_overrides.get("fps", scenario[FIELD_FPS]), name="fps" ), ) - return _run_replay_application( - app=T2VApplication(defaults=defaults), + return run_t2v_replay_application( + model=model_from_backend(defaults.backend, defaults.preset_id), + defaults=defaults, output=output, ) context = initialize_cuda_distributed(default_device=config.device) - adapter = make_adapter(defaults.backend) + adapter = make_adapter(defaults.backend, defaults.preset_id) output = WebRTCOutputSpec( host=str(host or output_overrides.get("host", "0.0.0.0")), port=_int_value( @@ -216,7 +159,6 @@ def launch_t2v( spec = _spec( replace(defaults, device=str(context.device)), adapter=adapter, - scenario=scenario, input_mode="webrtc", output=output, ) @@ -263,7 +205,11 @@ def create_app(config: "T2VDemoRunnerConfig | None" = None) -> Application: from .runner import RUNNER_T2V config = RUNNER_T2V if config is None else config - return T2VApplication(defaults=_defaults_from_config(config, {})) + defaults = _defaults_from_config(config, {}) + return create_t2v_application( + model=model_from_backend(defaults.backend, defaults.preset_id), + defaults=defaults, + ) createApp = create_app @@ -290,42 +236,24 @@ def value(name: str) -> object: def _scenario(defaults: T2VApplicationDefaults) -> dict[str, object]: - runner = resolve_backend(defaults.backend).resolve_runner(defaults.preset_id) - - def value(name: str, default: object) -> object: - configured = getattr(defaults, name) - return default if configured is None else configured - - return { - FIELD_PROMPT: value(FIELD_PROMPT, runner.prompt), - FIELD_TOTAL_BLOCKS: value(FIELD_TOTAL_BLOCKS, runner.total_blocks), - FIELD_PIXEL_HEIGHT: value(FIELD_PIXEL_HEIGHT, runner.pixel_height), - FIELD_PIXEL_WIDTH: value(FIELD_PIXEL_WIDTH, runner.pixel_width), - FIELD_FPS: value(FIELD_FPS, runner.fps), - } + return t2v_scenario_mapping( + model=model_from_backend(defaults.backend, defaults.preset_id), + defaults=defaults, + ) def _spec( defaults: T2VApplicationDefaults, *, adapter: T2VDemoAdapter, - scenario: dict[str, object], input_mode: Literal["replay", "webrtc"], output: Mp4OutputSpec | NullOutputSpec | WebRTCOutputSpec, ) -> DemoSpec: - return DemoSpec( - model_id=adapter.model_id, - preset_id=defaults.preset_id or adapter.backend.default_preset_name, + return create_t2v_spec( + model=adapter.model, + defaults=defaults, input_mode=input_mode, - scenario=scenario, output=output, - config=InferenceConfig( - model_id=adapter.model_id, - preset_id=defaults.preset_id or adapter.backend.default_preset_name, - device=defaults.device, - compile=defaults.compile, - runtime_options={"backend": adapter.backend.key}, - ), ) @@ -369,28 +297,6 @@ def _optional_str(value: object) -> str | None: return None if value is None else str(value) -def _run_replay_application( - *, - app: Application, - output: Mp4OutputSpec | NullOutputSpec, -) -> RunResult: - output_sink = None - if isinstance(output, Mp4OutputSpec): - output_sink = FileOutputSink( - output_path=output.path, - fps=output.fps, - output_layout=output.output_layout, - ) - result = Runner( - io_handler=create_replay_io_handler(output_sink=output_sink), - app=app, - ).run() - if result.status != "completed": - reason = result.reason or str(result.error) or "T2V replay failed." - raise RuntimeError(reason) - return result - - def _configure_app( app: web.Application, *, @@ -459,7 +365,6 @@ async def playback(_: web.Request) -> web.StreamResponse: __all__ = [ - "T2VApplication", "T2VApplicationDefaults", "T2VWebRTCSessionManager", "createApp", diff --git a/apps/t2v_demo/pyproject.toml b/apps/t2v_demo/pyproject.toml index 29698b039..e357fd44d 100644 --- a/apps/t2v_demo/pyproject.toml +++ b/apps/t2v_demo/pyproject.toml @@ -10,7 +10,7 @@ name = "flashdreams-t2v-demo" version = "0.1.0" description = "FlashDreams text-to-video runtime demo launcher" requires-python = ">=3.10" -dependencies = ["flashdreams[serving]"] +dependencies = ["flashdreams[serving]", "flashdreams-t2v"] [project.entry-points."flashdreams.runner_configs"] t2v = "t2v_demo.runner:RUNNER_T2V" @@ -20,6 +20,7 @@ t2v = "t2v_demo.app:create_app" [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [tool.setuptools] packages = ["t2v_demo"] diff --git a/apps/t2v_demo/runtime.py b/apps/t2v_demo/runtime.py index 873496c17..8e149ceba 100644 --- a/apps/t2v_demo/runtime.py +++ b/apps/t2v_demo/runtime.py @@ -1,297 +1,64 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime API adapter for the T2V demo's integration pipelines.""" +"""Compatibility bridge from the legacy T2V demo registry to the shared shell.""" from __future__ import annotations -from dataclasses import dataclass -from pathlib import Path -from typing import Any -from uuid import uuid4 - -import torch - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.video_output import VideoOutputStream -from flashdreams.runtime import ( - CanonicalInputSchema, - IdentityInputMapping, - InferenceConfig, - InferenceInput, - InferenceInputSchema, - InputField, - ModelAdapter, - StepRequest, -) -from flashdreams.runtime.demo import DemoSpec, PreparedScenario -from flashdreams.runtime.demo.session_inputs import ( - PreparedStep, - ProviderCapabilities, - UserInputWindow, +from t2v.t2v import ( + FIELD_FPS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + T2VDemoAdapter, + T2VInputProvider, + T2VModelConfig, + T2VRuntime, + T2VScenario, + T2VSession, ) -from flashdreams.runtime.interfaces import InferenceSession -from flashdreams.runtime.types import StepResult -from flashdreams.runtime.video_output import Mp4VideoOutputTarget from .backends import T2VBackend, resolve_backend -FIELD_PROMPT = "prompt" -FIELD_TOTAL_BLOCKS = "total_blocks" -FIELD_PIXEL_HEIGHT = "pixel_height" -FIELD_PIXEL_WIDTH = "pixel_width" -FIELD_FPS = "fps" - - -@dataclass(frozen=True, kw_only=True, slots=True) -class T2VScenario: - """Prompt and output geometry for a finite text-to-video rollout.""" - - prompt: str - total_blocks: int - pixel_height: int - pixel_width: int - fps: int - - -class T2VDemoAdapter(ModelAdapter): - """Model adapter shared by replay and WebRTC T2V launch paths.""" - - model_id = "flashdreams-t2v" - inference_input_schema = InferenceInputSchema( - global_conditioning_fields=(InputField(name=FIELD_PROMPT),), - description="Text-to-video prompt and rollout settings.", - ) - canonical_input_schema = CanonicalInputSchema() - - def __init__(self, *, backend: T2VBackend) -> None: - self.backend = backend - - def supported_input_modes(self) -> tuple[str, ...]: - return ("replay", "webrtc") - - def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4", "null", "webrtc") - - def default_input_mapping(self) -> IdentityInputMapping: - return IdentityInputMapping() - - def validate_config(self, config: InferenceConfig) -> None: - if config.model_id != self.model_id: - raise ValueError( - f"Expected model_id={self.model_id!r}, got {config.model_id!r}." - ) - if config.runtime_options.get("backend") != self.backend.key: - raise ValueError("T2V runtime backend does not match its demo adapter.") - - def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: - scenario = _scenario_from_value(spec.scenario, self.backend) - return PreparedScenario( - initial_inputs=InferenceInput( - global_conditioning={ - FIELD_PROMPT: scenario.prompt, - FIELD_TOTAL_BLOCKS: scenario.total_blocks, - FIELD_PIXEL_HEIGHT: scenario.pixel_height, - FIELD_PIXEL_WIDTH: scenario.pixel_width, - FIELD_FPS: scenario.fps, - } - ) - ) - - def create_runtime(self, config: InferenceConfig) -> "T2VRuntime": - self.validate_config(config) - return T2VRuntime(config=config, backend=self.backend) - - def create_model_input_provider( - self, spec: DemoSpec, scenario: PreparedScenario - ) -> "T2VInputProvider": - """Supply fixed prompt conditioning to every shared-demo step.""" - del spec - return T2VInputProvider(initial_inputs=scenario.initial_inputs) - - -class T2VInputProvider: - """No-control input provider for finite prompt-only generation.""" - - capabilities = ProviderCapabilities( - supports_realtime_clock=True, - supports_recorded_input=True, - deterministic_given_inputs=True, - ) - - def __init__(self, *, initial_inputs: InferenceInput) -> None: - self._initial_inputs = initial_inputs - - def prepare_initial_input(self) -> InferenceInput: - return self._initial_inputs - - def prepare_step( - self, *, request: Any, user_window: UserInputWindow - ) -> PreparedStep: - del request, user_window - return PreparedStep(inference_input=InferenceInput()) - - def reset(self, inputs: InferenceInput | None = None) -> None: - if inputs is not None: - self._initial_inputs = inputs - - def close(self) -> None: - pass - - -class T2VRuntime: - """One heavyweight selected pipeline, reusable across demo sessions.""" - - def __init__(self, *, config: InferenceConfig, backend: T2VBackend) -> None: - self.config = config - self.backend = backend - runner = backend.resolve_runner(config.preset_id) - pipeline_config = runner.pipeline - if config.compile is not None: - from flashdreams.infra.config import derive_config - - pipeline_config = derive_config( - base_config=pipeline_config, - diffusion_model={"transformer": {"compile_network": config.compile}}, - ) - self.pipeline = pipeline_config.setup().to(config.device or "cuda").eval() - self._latest_artifact: tuple[Path, T2VScenario] | None = None - - def blocks_for_duration(self, duration_s: float, *, fps: int) -> int: - """Return enough autoregressive chunks to reach the requested duration.""" - target_frames = int(duration_s * fps) - frames = 0 - index = 0 - while frames < target_frames: - frames += int(self.pipeline.get_num_output_frames(index)) - index += 1 - return index - - def record_artifact(self, path: Path, scenario: T2VScenario) -> None: - self._latest_artifact = (path, scenario) - - @property - def latest_artifact(self) -> tuple[Path, T2VScenario] | None: - return self._latest_artifact - - def start_session(self, inputs: InferenceInput) -> "T2VSession": - return T2VSession( - pipeline=self.pipeline, scenario=_scenario_from_inputs(inputs), runtime=self - ) - - def close(self) -> None: - close = getattr(self.pipeline, "close", None) - if callable(close): - close() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - -class T2VSession(InferenceSession): - """A cache-isolated T2V session that yields chunks as they are generated.""" - - def __init__( - self, *, pipeline: Any, scenario: T2VScenario, runtime: T2VRuntime - ) -> None: - self.pipeline = pipeline - self.scenario = scenario - self._runtime = runtime - self._artifact_path = Path("outputs/t2v-webrtc") / f"{uuid4()}.mp4" - self._artifact_path.parent.mkdir(parents=True, exist_ok=True) - self._artifact_output = Mp4VideoOutputTarget( - output_path=self._artifact_path, fps=scenario.fps, output_layout="tchw" - ) - self._artifact_output.open() - self._step_index = 0 - self._closed = False - self._output_stream = VideoOutputStream( - postprocess_stream=None, output_layout="tchw" - ) - assert isinstance(pipeline.decoder, StreamingVideoDecoder) - ratio = pipeline.decoder.spatial_compression_ratio - if scenario.pixel_height % ratio or scenario.pixel_width % ratio: - raise ValueError( - "T2V dimensions must be divisible by the decoder spatial compression ratio." - ) - self._cache = pipeline.initialize_cache( - text=[scenario.prompt], - image=None, - height=scenario.pixel_height // ratio, - width=scenario.pixel_width // ratio, - ) - - def next_step_request(self) -> StepRequest | None: - if self._closed or self._step_index >= self.scenario.total_blocks: - return None - return StepRequest(step_index=self._step_index) - - def step(self, inputs: InferenceInput) -> StepResult: - del inputs - if self._closed: - raise RuntimeError("T2V session is closed.") - index = self._step_index - video = self.pipeline.generate(autoregressive_index=index, cache=self._cache) - stats = self.pipeline.finalize(autoregressive_index=index, cache=self._cache) - self._step_index += 1 - result = self._output_stream.process( - video, - autoregressive_index=index, - metrics=stats, - metadata={"prompt": self.scenario.prompt}, - ) - self._artifact_output.write(result) - return result - - def reset(self, inputs: InferenceInput | None = None) -> None: - if inputs is not None and _scenario_from_inputs(inputs) != self.scenario: - raise ValueError( - "Create a new T2V session to change the prompt or dimensions." - ) - raise RuntimeError( - "T2V sessions are finite; create a new session instead of reset()." - ) - - def close(self) -> None: - if self._closed: - return - self._closed = True - artifacts = self._artifact_output.close() - if artifacts: - self._runtime.record_artifact(self._artifact_path, self.scenario) - - -def _scenario_from_value(value: Any, backend: T2VBackend) -> T2VScenario: - runner = backend.resolve_runner() - source = value if isinstance(value, dict) else {} - prompt = str(source.get(FIELD_PROMPT, getattr(runner, FIELD_PROMPT, ""))).strip() - if not prompt: - raise ValueError("A non-empty text-to-video prompt is required.") - return T2VScenario( - prompt=prompt, - total_blocks=int( - source.get(FIELD_TOTAL_BLOCKS, getattr(runner, FIELD_TOTAL_BLOCKS, 1)) - ), - pixel_height=int( - source.get(FIELD_PIXEL_HEIGHT, getattr(runner, FIELD_PIXEL_HEIGHT, 480)) - ), - pixel_width=int( - source.get(FIELD_PIXEL_WIDTH, getattr(runner, FIELD_PIXEL_WIDTH, 832)) - ), - fps=int(source.get(FIELD_FPS, getattr(runner, FIELD_FPS, 16))), - ) - -def _scenario_from_inputs(inputs: InferenceInput) -> T2VScenario: - source = inputs.global_conditioning - return T2VScenario( - prompt=str(source[FIELD_PROMPT]), - total_blocks=int(source[FIELD_TOTAL_BLOCKS]), - pixel_height=int(source[FIELD_PIXEL_HEIGHT]), - pixel_width=int(source[FIELD_PIXEL_WIDTH]), - fps=int(source[FIELD_FPS]), +def model_from_backend( + backend: str | T2VBackend, + preset_id: str | None = None, +) -> T2VModelConfig: + """Adapt the existing backend registry to one neutral T2V model config.""" + resolved = resolve_backend(backend) if isinstance(backend, str) else backend + preset = resolved.resolve_runner(preset_id) + return T2VModelConfig( + model_id="flashdreams-t2v", + preset_id=preset.name, + pipeline=preset.pipeline, + prompt=preset.prompt, + total_blocks=preset.total_blocks, + pixel_height=preset.pixel_height, + pixel_width=preset.pixel_width, + fps=preset.fps, + runtime_options={"backend": resolved.key}, ) -def make_adapter(backend: str) -> T2VDemoAdapter: +def make_adapter(backend: str, preset_id: str | None = None) -> T2VDemoAdapter: """Build an adapter from a CLI/UI backend key.""" - return T2VDemoAdapter(backend=resolve_backend(backend)) + return T2VDemoAdapter(model=model_from_backend(backend, preset_id)) + + +__all__ = [ + "FIELD_FPS", + "FIELD_PIXEL_HEIGHT", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "T2VDemoAdapter", + "T2VInputProvider", + "T2VModelConfig", + "T2VRuntime", + "T2VScenario", + "T2VSession", + "make_adapter", + "model_from_backend", +] diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py index 47e945604..fc4d928c6 100644 --- a/apps/t2v_demo/tests/test_runner.py +++ b/apps/t2v_demo/tests/test_runner.py @@ -6,11 +6,18 @@ from pathlib import Path import pytest +import t2v.t2v as t2v_shell import tomli from t2v_demo import app from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig - -from flashdreams.demo import Application, FileOutputSink, ReplayIOHandler +from t2v_demo.runtime import model_from_backend + +from flashdreams.demo import ( + Application, + DemoAdapterApplication, + FileOutputSink, + ReplayIOHandler, +) from flashdreams.runtime.demo import RunResult pytestmark = pytest.mark.ci_cpu @@ -43,10 +50,23 @@ def test_t2v_create_app_exposes_public_application() -> None: assert app.createApp is app.create_app assert isinstance(public_app, Application) - assert isinstance(public_app, app.T2VApplication) - assert public_app.defaults.backend == "self-forcing" - assert public_app.defaults.prompt == "A waterfall" - assert public_app.defaults.total_blocks == 3 + assert isinstance(public_app, DemoAdapterApplication) + spec = public_app.spec + assert spec.model_id == "flashdreams-t2v" + assert spec.config is not None + assert spec.config.runtime_options["backend"] == "self-forcing" + scenario = spec.scenario + assert isinstance(scenario, dict) + assert scenario["prompt"] == "A waterfall" + assert scenario["total_blocks"] == 3 + + +def test_t2v_backend_bridge_builds_neutral_model_config() -> None: + model = model_from_backend("self-forcing") + + assert isinstance(model, t2v_shell.T2VModelConfig) + assert model.model_id == "flashdreams-t2v" + assert model.runtime_options["backend"] == "self-forcing" def test_runner_mp4_launch_uses_demo_entrypoint( @@ -62,7 +82,7 @@ def __init__(self, *, io_handler: object, app: object) -> None: def run(self) -> RunResult: return RunResult(status="completed") - monkeypatch.setattr(app, "Runner", FakeRunner) + monkeypatch.setattr(t2v_shell, "Runner", FakeRunner) config = T2VDemoRunnerConfig( runner_name="t2v", description="test", @@ -79,10 +99,12 @@ def run(self) -> RunResult: public_app = captured["app"] io_handler = captured["io_handler"] - assert isinstance(public_app, app.T2VApplication) + assert isinstance(public_app, DemoAdapterApplication) assert isinstance(io_handler, ReplayIOHandler) - assert public_app.defaults.prompt == "A waterfall" - assert public_app.defaults.total_blocks == 3 + scenario = public_app.spec.scenario + assert isinstance(scenario, dict) + assert scenario["prompt"] == "A waterfall" + assert scenario["total_blocks"] == 3 output_sink = io_handler.output_sink assert isinstance(output_sink, FileOutputSink) assert str(output_sink.output_path) == "outputs/test.mp4" diff --git a/uv.lock b/uv.lock index 7f4aba302..b5b722dee 100644 --- a/uv.lock +++ b/uv.lock @@ -29,6 +29,7 @@ members = [ "flashdreams-omnidreams", "flashdreams-sana-wm", "flashdreams-self-forcing", + "flashdreams-t2v", "flashdreams-t2v-demo", "flashdreams-wan21", "flashdreams-wan22", @@ -1418,16 +1419,31 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "flashdreams-t2v" +version = "0.1.0" +source = { editable = "apps/t2v" } +dependencies = [ + { name = "flashdreams" }, +] + +[package.metadata] +requires-dist = [{ name = "flashdreams", editable = "flashdreams" }] + [[package]] name = "flashdreams-t2v-demo" version = "0.1.0" source = { editable = "apps/t2v_demo" } dependencies = [ { name = "flashdreams", extra = ["serving"] }, + { name = "flashdreams-t2v" }, ] [package.metadata] -requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] +requires-dist = [ + { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, +] [[package]] name = "flashdreams-wan21" From c571b86a3bf340949e1701c7f1afc015abb7b973 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 10:33:34 +0000 Subject: [PATCH 20/33] Add integration-owned T2V app entries --- apps/t2v/t2v.py | 7 +-- apps/t2v/tests/test_t2v_shell.py | 6 +-- .../causal_forcing/t2v/__init__.py | 8 ++++ .../causal_forcing/causal_forcing/t2v/app.py | 31 +++++++++++++ integrations/causal_forcing/pyproject.toml | 5 +++ .../causal_forcing/tests/test_smoke.py | 29 ++++++++++++- .../cosmos_predict2/t2v/__init__.py | 8 ++++ .../cosmos_predict2/t2v/app.py | 31 +++++++++++++ integrations/cosmos_predict2/pyproject.toml | 5 +++ .../cosmos_predict2/tests/test_t2v_app.py | 43 +++++++++++++++++++ integrations/self_forcing/pyproject.toml | 5 +++ .../self_forcing/self_forcing/t2v/__init__.py | 8 ++++ .../self_forcing/self_forcing/t2v/app.py | 31 +++++++++++++ integrations/self_forcing/tests/test_smoke.py | 29 ++++++++++++- uv.lock | 6 +++ 15 files changed, 244 insertions(+), 8 deletions(-) create mode 100644 integrations/causal_forcing/causal_forcing/t2v/__init__.py create mode 100644 integrations/causal_forcing/causal_forcing/t2v/app.py create mode 100644 integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py create mode 100644 integrations/cosmos_predict2/cosmos_predict2/t2v/app.py create mode 100644 integrations/cosmos_predict2/tests/test_t2v_app.py create mode 100644 integrations/self_forcing/self_forcing/t2v/__init__.py create mode 100644 integrations/self_forcing/self_forcing/t2v/app.py diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index 585b3cbf9..64a1d017a 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -154,8 +154,7 @@ def default_input_mapping(self) -> IdentityInputMapping: def validate_config(self, config: InferenceConfig) -> None: if config.model_id != self.model.model_id: raise ValueError( - f"Expected model_id={self.model.model_id!r}, " - f"got {config.model_id!r}." + f"Expected model_id={self.model.model_id!r}, got {config.model_id!r}." ) if ( self.model.preset_id is not None @@ -490,7 +489,9 @@ def _int_value(value: object) -> int: return int(value) if isinstance(value, str): return int(value) - raise TypeError(f"Expected an integer-compatible value, got {type(value).__name__}.") + raise TypeError( + f"Expected an integer-compatible value, got {type(value).__name__}." + ) def _validate_positive_int(value: int, *, name: str) -> None: diff --git a/apps/t2v/tests/test_t2v_shell.py b/apps/t2v/tests/test_t2v_shell.py index d0bd02079..f0202e51b 100644 --- a/apps/t2v/tests/test_t2v_shell.py +++ b/apps/t2v/tests/test_t2v_shell.py @@ -6,9 +6,6 @@ import inspect import pytest - -from flashdreams.demo import DemoAdapterApplication -from flashdreams.runtime.demo import Mp4OutputSpec, NullOutputSpec from t2v.t2v import ( T2VDemoAdapter, T2VModelConfig, @@ -18,6 +15,9 @@ t2v_scenario_mapping, ) +from flashdreams.demo import DemoAdapterApplication +from flashdreams.runtime.demo import Mp4OutputSpec, NullOutputSpec + pytestmark = pytest.mark.ci_cpu diff --git a/integrations/causal_forcing/causal_forcing/t2v/__init__.py b/integrations/causal_forcing/causal_forcing/t2v/__init__.py new file mode 100644 index 000000000..175fd535e --- /dev/null +++ b/integrations/causal_forcing/causal_forcing/t2v/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Causal-Forcing T2V public demo app.""" + +from causal_forcing.t2v.app import MODEL, createApp, create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/causal_forcing/causal_forcing/t2v/app.py b/integrations/causal_forcing/causal_forcing/t2v/app.py new file mode 100644 index 000000000..ab71a6d31 --- /dev/null +++ b/integrations/causal_forcing/causal_forcing/t2v/app.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public T2V app entry for the default Causal-Forcing model.""" + +from flashdreams.demo import Application +from t2v import T2VModelConfig, create_t2v_application + +from causal_forcing.config import PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE +from causal_forcing.runner import DEFAULT_T2V_PROMPT + +MODEL = T2VModelConfig( + model_id="causal-forcing-t2v", + preset_id=PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE.name, + pipeline=PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE, + prompt=DEFAULT_T2V_PROMPT, + total_blocks=60, + pixel_height=480, + pixel_width=832, + fps=16, +) + + +def create_app() -> Application: + """Create the default Causal-Forcing T2V application.""" + return create_t2v_application(model=MODEL) + + +createApp = create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/causal_forcing/pyproject.toml b/integrations/causal_forcing/pyproject.toml index 1105845b9..76425f307 100644 --- a/integrations/causal_forcing/pyproject.toml +++ b/integrations/causal_forcing/pyproject.toml @@ -25,12 +25,14 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "flashdreams", + "flashdreams-t2v", "mediapy>=1.1", "opencv-python-headless>=4.5", ] [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [project.optional-dependencies] dev = [ @@ -46,6 +48,9 @@ dev = [ "causal-forcing-wan2.1-t2v-1.3b-framewise" = "causal_forcing.config:RUNNER_WAN21_T2V_1PT3B_FRAMEWISE" "causal-forcing-wan2.1-i2v-1.3b-framewise" = "causal_forcing.config:RUNNER_WAN21_I2V_1PT3B_FRAMEWISE" +[project.entry-points."flashdreams.applications"] +causal-forcing-t2v = "causal_forcing.t2v.app:create_app" + [tool.setuptools.packages.find] include = ["causal_forcing*"] exclude = ["tests"] diff --git a/integrations/causal_forcing/tests/test_smoke.py b/integrations/causal_forcing/tests/test_smoke.py index 3ef01f2e5..39c5ff9d4 100644 --- a/integrations/causal_forcing/tests/test_smoke.py +++ b/integrations/causal_forcing/tests/test_smoke.py @@ -32,13 +32,16 @@ import pytest import tomli as tomllib from causal_forcing import config as config_mod -from causal_forcing.config import RUNNER_CONFIGS +from causal_forcing.config import PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE, RUNNER_CONFIGS +from causal_forcing.t2v.app import MODEL, create_app, createApp +from flashdreams.demo import Application, DemoAdapterApplication from flashdreams.infra.runner import RunnerConfig pytestmark = pytest.mark.ci_gpu ENTRY_POINT_GROUP = "flashdreams.runner_configs" +APPLICATION_ENTRY_POINT_GROUP = "flashdreams.applications" def test_runners_dict_is_non_empty() -> None: @@ -97,6 +100,30 @@ def test_entry_points_match_module_literals() -> None: ) +def test_t2v_app_uses_default_pipeline_config() -> None: + """The public app entry must remain owned by this integration package.""" + public_app = create_app() + + assert createApp is create_app + assert isinstance(public_app, Application) + assert isinstance(public_app, DemoAdapterApplication) + assert MODEL.model_id == "causal-forcing-t2v" + assert MODEL.preset_id == PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE.name + assert MODEL.pipeline is PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE + assert public_app.spec.model_id == MODEL.model_id + assert public_app.spec.preset_id == MODEL.preset_id + + +def test_application_entry_point_matches_module_literal() -> None: + """The integration owns its public T2V application entry point.""" + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as fh: + meta = tomllib.load(fh) + + entries = meta["project"]["entry-points"][APPLICATION_ENTRY_POINT_GROUP] + assert entries == {"causal-forcing-t2v": "causal_forcing.t2v.app:create_app"} + + @pytest.mark.skipif( sys.version_info < (3, 11), reason="entry-point discovery test relies on ``importlib.metadata`` 3.10+ shape", diff --git a/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py b/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py new file mode 100644 index 000000000..a72ebbf6e --- /dev/null +++ b/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cosmos Predict2 T2V public demo app.""" + +from cosmos_predict2.t2v.app import MODEL, createApp, create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py b/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py new file mode 100644 index 000000000..a9467fc4b --- /dev/null +++ b/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public T2V app entry for the default Cosmos Predict2 model.""" + +from flashdreams.demo import Application +from t2v import T2VModelConfig, create_t2v_application + +from cosmos_predict2.config import PIPELINE_COSMOS2_T2V_2B_720P +from cosmos_predict2.runner import DEFAULT_PROMPT + +MODEL = T2VModelConfig( + model_id="cosmos-predict2-t2v", + preset_id=PIPELINE_COSMOS2_T2V_2B_720P.name, + pipeline=PIPELINE_COSMOS2_T2V_2B_720P, + prompt=DEFAULT_PROMPT, + total_blocks=1, + pixel_height=720, + pixel_width=1280, + fps=16, +) + + +def create_app() -> Application: + """Create the default Cosmos Predict2 T2V application.""" + return create_t2v_application(model=MODEL) + + +createApp = create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/cosmos_predict2/pyproject.toml b/integrations/cosmos_predict2/pyproject.toml index ef92ec735..dd429a37c 100644 --- a/integrations/cosmos_predict2/pyproject.toml +++ b/integrations/cosmos_predict2/pyproject.toml @@ -25,11 +25,13 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "flashdreams", + "flashdreams-t2v", "mediapy>=1.1", ] [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [project.optional-dependencies] dev = [ @@ -44,6 +46,9 @@ dev = [ "cosmos2-t2v-2b-720p" = "cosmos_predict2.config:RUNNER_COSMOS2_T2V_2B_720P" "cosmos2-i2v-2b-720p" = "cosmos_predict2.config:RUNNER_COSMOS2_I2V_2B_720P" +[project.entry-points."flashdreams.applications"] +cosmos-predict2-t2v = "cosmos_predict2.t2v.app:create_app" + [tool.setuptools.packages.find] include = ["cosmos_predict2*"] exclude = ["tests"] diff --git a/integrations/cosmos_predict2/tests/test_t2v_app.py b/integrations/cosmos_predict2/tests/test_t2v_app.py new file mode 100644 index 000000000..fa3f3483e --- /dev/null +++ b/integrations/cosmos_predict2/tests/test_t2v_app.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import pytest +import tomli as tomllib +from cosmos_predict2.config import PIPELINE_COSMOS2_T2V_2B_720P +from cosmos_predict2.t2v.app import MODEL, createApp, create_app + +from flashdreams.demo import Application, DemoAdapterApplication + +pytestmark = pytest.mark.ci_cpu + +APPLICATION_ENTRY_POINT_GROUP = "flashdreams.applications" + + +def test_t2v_app_uses_default_pipeline_config() -> None: + """The public app entry must remain owned by this integration package.""" + public_app = create_app() + + assert createApp is create_app + assert isinstance(public_app, Application) + assert isinstance(public_app, DemoAdapterApplication) + assert MODEL.model_id == "cosmos-predict2-t2v" + assert MODEL.preset_id == PIPELINE_COSMOS2_T2V_2B_720P.name + assert MODEL.pipeline is PIPELINE_COSMOS2_T2V_2B_720P + assert public_app.spec.model_id == MODEL.model_id + assert public_app.spec.preset_id == MODEL.preset_id + + +def test_application_entry_point_matches_module_literal() -> None: + """The integration owns its public T2V application entry point.""" + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as fh: + meta = tomllib.load(fh) + + entries = meta["project"]["entry-points"][APPLICATION_ENTRY_POINT_GROUP] + assert entries == { + "cosmos-predict2-t2v": "cosmos_predict2.t2v.app:create_app" + } diff --git a/integrations/self_forcing/pyproject.toml b/integrations/self_forcing/pyproject.toml index abd0f9260..d8d28bd3a 100644 --- a/integrations/self_forcing/pyproject.toml +++ b/integrations/self_forcing/pyproject.toml @@ -25,11 +25,13 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "flashdreams", + "flashdreams-t2v", "mediapy>=1.1", ] [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [project.optional-dependencies] dev = [ @@ -45,6 +47,9 @@ dev = [ "self-forcing-wan2.1-t2v-1.3b-taehv" = "self_forcing.config:RUNNER_WAN21_T2V_1PT3B_TAEHV" "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope" = "self_forcing.config:RUNNER_WAN21_T2V_1PT3B_SINK5_WINDOW7_REROPE" +[project.entry-points."flashdreams.applications"] +self-forcing-t2v = "self_forcing.t2v.app:create_app" + [tool.setuptools.packages.find] include = ["self_forcing*"] exclude = ["tests"] diff --git a/integrations/self_forcing/self_forcing/t2v/__init__.py b/integrations/self_forcing/self_forcing/t2v/__init__.py new file mode 100644 index 000000000..1ed51e404 --- /dev/null +++ b/integrations/self_forcing/self_forcing/t2v/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Self-Forcing T2V public demo app.""" + +from self_forcing.t2v.app import MODEL, createApp, create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/self_forcing/self_forcing/t2v/app.py b/integrations/self_forcing/self_forcing/t2v/app.py new file mode 100644 index 000000000..2eb64cf35 --- /dev/null +++ b/integrations/self_forcing/self_forcing/t2v/app.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public T2V app entry for the default Self-Forcing model.""" + +from flashdreams.demo import Application +from t2v import T2VModelConfig, create_t2v_application + +from self_forcing.config import PIPELINE_WAN21_T2V_1PT3B +from self_forcing.runner import DEFAULT_T2V_PROMPT + +MODEL = T2VModelConfig( + model_id="self-forcing-t2v", + preset_id=PIPELINE_WAN21_T2V_1PT3B.name, + pipeline=PIPELINE_WAN21_T2V_1PT3B, + prompt=DEFAULT_T2V_PROMPT, + total_blocks=60, + pixel_height=480, + pixel_width=832, + fps=16, +) + + +def create_app() -> Application: + """Create the default Self-Forcing T2V application.""" + return create_t2v_application(model=MODEL) + + +createApp = create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/self_forcing/tests/test_smoke.py b/integrations/self_forcing/tests/test_smoke.py index 77de71b5f..3e6377fe9 100644 --- a/integrations/self_forcing/tests/test_smoke.py +++ b/integrations/self_forcing/tests/test_smoke.py @@ -24,13 +24,16 @@ import pytest import tomli as tomllib from self_forcing import config as config_mod -from self_forcing.config import RUNNER_CONFIGS +from self_forcing.config import PIPELINE_WAN21_T2V_1PT3B, RUNNER_CONFIGS +from self_forcing.t2v.app import MODEL, create_app, createApp +from flashdreams.demo import Application, DemoAdapterApplication from flashdreams.infra.runner import RunnerConfig pytestmark = pytest.mark.ci_cpu ENTRY_POINT_GROUP = "flashdreams.runner_configs" +APPLICATION_ENTRY_POINT_GROUP = "flashdreams.applications" def test_runners_dict_is_non_empty() -> None: @@ -89,6 +92,30 @@ def test_entry_points_match_module_literals() -> None: ) +def test_t2v_app_uses_default_pipeline_config() -> None: + """The public app entry must remain owned by this integration package.""" + public_app = create_app() + + assert createApp is create_app + assert isinstance(public_app, Application) + assert isinstance(public_app, DemoAdapterApplication) + assert MODEL.model_id == "self-forcing-t2v" + assert MODEL.preset_id == PIPELINE_WAN21_T2V_1PT3B.name + assert MODEL.pipeline is PIPELINE_WAN21_T2V_1PT3B + assert public_app.spec.model_id == MODEL.model_id + assert public_app.spec.preset_id == MODEL.preset_id + + +def test_application_entry_point_matches_module_literal() -> None: + """The integration owns its public T2V application entry point.""" + pyproject = Path(__file__).resolve().parents[1] / "pyproject.toml" + with pyproject.open("rb") as fh: + meta = tomllib.load(fh) + + entries = meta["project"]["entry-points"][APPLICATION_ENTRY_POINT_GROUP] + assert entries == {"self-forcing-t2v": "self_forcing.t2v.app:create_app"} + + @pytest.mark.skipif( sys.version_info < (3, 11), reason="entry-point discovery test relies on ``importlib.metadata`` 3.10+ shape", diff --git a/uv.lock b/uv.lock index b5b722dee..f8190bfd3 100644 --- a/uv.lock +++ b/uv.lock @@ -1114,6 +1114,7 @@ version = "0.1.0" source = { editable = "integrations/causal_forcing" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, { name = "opencv-python-headless" }, ] @@ -1126,6 +1127,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, { name = "mediapy", specifier = ">=1.1" }, { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, @@ -1138,6 +1140,7 @@ version = "0.1.0" source = { editable = "integrations/cosmos_predict2" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1149,6 +1152,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, { name = "mediapy", specifier = ">=1.1" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, ] @@ -1403,6 +1407,7 @@ version = "0.1.0" source = { editable = "integrations/self_forcing" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1414,6 +1419,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, { name = "mediapy", specifier = ">=1.1" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, ] From 1df1e7e45b56902d2a8291382a353f6cc1de49d9 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 10:44:14 +0000 Subject: [PATCH 21/33] Remove duplicated T2V demo preset registry Derive legacy t2v_demo backend metadata and preset selection from the integration-owned application and runner configs. Add a neutral helper for building T2V model configs from runner configs, and keep the old t2v runner bridge working while deleting the app-local backends.py and presets.py copies. --- apps/t2v/__init__.py | 2 + apps/t2v/t2v.py | 21 ++ apps/t2v/tests/test_t2v_shell.py | 27 ++ apps/t2v_demo/app.py | 2 +- apps/t2v_demo/backends.py | 88 ------- apps/t2v_demo/presets.py | 243 ------------------ apps/t2v_demo/pyproject.toml | 11 +- apps/t2v_demo/runner.py | 6 +- apps/t2v_demo/runtime.py | 190 ++++++++++++-- apps/t2v_demo/tests/test_runner.py | 35 ++- .../causal_forcing/t2v/__init__.py | 2 +- .../causal_forcing/causal_forcing/t2v/app.py | 17 +- .../causal_forcing/tests/test_smoke.py | 6 +- .../cosmos_predict2/t2v/__init__.py | 2 +- .../cosmos_predict2/t2v/app.py | 17 +- .../cosmos_predict2/tests/test_t2v_app.py | 12 +- .../self_forcing/self_forcing/t2v/__init__.py | 2 +- .../self_forcing/self_forcing/t2v/app.py | 17 +- integrations/self_forcing/tests/test_smoke.py | 6 +- uv.lock | 6 + 20 files changed, 302 insertions(+), 410 deletions(-) delete mode 100644 apps/t2v_demo/backends.py delete mode 100644 apps/t2v_demo/presets.py diff --git a/apps/t2v/__init__.py b/apps/t2v/__init__.py index 920e1595c..6bde43ec6 100644 --- a/apps/t2v/__init__.py +++ b/apps/t2v/__init__.py @@ -18,6 +18,7 @@ T2VSession, create_t2v_application, create_t2v_spec, + model_config_from_runner, run_t2v_replay_application, t2v_scenario_mapping, ) @@ -37,6 +38,7 @@ "T2VSession", "create_t2v_application", "create_t2v_spec", + "model_config_from_runner", "run_t2v_replay_application", "t2v_scenario_mapping", ] diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index 64a1d017a..583ea8b15 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -379,6 +379,26 @@ def create_t2v_application( return DemoAdapterApplication(adapter=T2VDemoAdapter(model=model), spec=spec) +def model_config_from_runner( + *, + model_id: str, + runner: Any, + runtime_options: Mapping[str, Any] | None = None, +) -> T2VModelConfig: + """Create a T2V model config from an integration-owned runner config.""" + return T2VModelConfig( + model_id=model_id, + preset_id=str(runner.runner_name), + pipeline=runner.pipeline, + prompt=str(getattr(runner, FIELD_PROMPT)), + total_blocks=_int_value(getattr(runner, FIELD_TOTAL_BLOCKS, 1)), + pixel_height=_int_value(getattr(runner, FIELD_PIXEL_HEIGHT)), + pixel_width=_int_value(getattr(runner, FIELD_PIXEL_WIDTH)), + fps=_int_value(getattr(runner, FIELD_FPS)), + runtime_options=runtime_options or {}, + ) + + def create_t2v_spec( *, model: T2VModelConfig, @@ -522,6 +542,7 @@ def _validate_optional_positive_int(value: int | None, *, name: str) -> None: "T2VSession", "create_t2v_application", "create_t2v_spec", + "model_config_from_runner", "run_t2v_replay_application", "t2v_scenario_mapping", ] diff --git a/apps/t2v/tests/test_t2v_shell.py b/apps/t2v/tests/test_t2v_shell.py index f0202e51b..c3e4879e1 100644 --- a/apps/t2v/tests/test_t2v_shell.py +++ b/apps/t2v/tests/test_t2v_shell.py @@ -4,6 +4,7 @@ from __future__ import annotations import inspect +from types import SimpleNamespace import pytest from t2v.t2v import ( @@ -12,6 +13,7 @@ T2VRunDefaults, create_t2v_application, create_t2v_spec, + model_config_from_runner, t2v_scenario_mapping, ) @@ -70,6 +72,31 @@ def test_t2v_shell_creates_demo_adapter_application() -> None: assert public_app.spec.output.mode == "mp4" +def test_t2v_shell_builds_model_config_from_runner_config() -> None: + runner = SimpleNamespace( + runner_name="fake-runner", + pipeline=object(), + prompt="Runner prompt", + total_blocks=5, + pixel_height=48, + pixel_width=96, + fps=12, + ) + + model = model_config_from_runner( + model_id="fake-t2v", + runner=runner, + runtime_options={"owner": "integration"}, + ) + + assert model.model_id == "fake-t2v" + assert model.preset_id == "fake-runner" + assert model.pipeline is runner.pipeline + assert model.prompt == "Runner prompt" + assert model.total_blocks == 5 + assert model.runtime_options["owner"] == "integration" + + def test_t2v_shell_has_no_legacy_backend_imports() -> None: import t2v.t2v as t2v_shell diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index e29ec0a51..438d16d19 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -39,7 +39,6 @@ from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig -from .backends import backend_metadata from .runtime import ( FIELD_FPS, FIELD_PIXEL_HEIGHT, @@ -47,6 +46,7 @@ FIELD_PROMPT, FIELD_TOTAL_BLOCKS, T2VDemoAdapter, + backend_metadata, make_adapter, model_from_backend, ) diff --git a/apps/t2v_demo/backends.py b/apps/t2v_demo/backends.py deleted file mode 100644 index da2c7f0b9..000000000 --- a/apps/t2v_demo/backends.py +++ /dev/null @@ -1,88 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Text-to-video backend selection for the unified demo.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from .presets import PRESETS, T2VPreset - - -@dataclass(frozen=True, slots=True) -class T2VBackend: - """One supported T2V model family and its app-owned presets.""" - - key: str - label: str - default_preset_name: str - preset_names: tuple[str, ...] - - def resolve_runner(self, preset_name: str | None = None) -> T2VPreset: - """Return a preset without importing an integration package.""" - name = preset_name or self.default_preset_name - if name not in self.preset_names: - raise ValueError( - f"Unknown {self.key} preset {name!r}. Available presets: " - f"{', '.join(self.preset_names)}." - ) - return PRESETS[name] - - -BACKENDS: dict[str, T2VBackend] = { - "causal-forcing": T2VBackend( - key="causal-forcing", - label="Causal-Forcing (Wan 2.1)", - default_preset_name="causal-forcing-wan2.1-t2v-1.3b-chunkwise", - preset_names=( - "causal-forcing-wan2.1-t2v-1.3b-chunkwise", - "causal-forcing-wan2.1-t2v-1.3b-framewise", - ), - ), - "cosmos-predict2": T2VBackend( - key="cosmos-predict2", - label="Cosmos Predict2", - default_preset_name="cosmos2-t2v-2b-720p", - preset_names=("cosmos2-t2v-2b-720p",), - ), - "self-forcing": T2VBackend( - key="self-forcing", - label="Self-Forcing (Wan 2.1)", - default_preset_name="self-forcing-wan2.1-t2v-1.3b", - preset_names=( - "self-forcing-wan2.1-t2v-1.3b", - "self-forcing-wan2.1-t2v-1.3b-taehv", - "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope", - ), - ), -} - - -def resolve_backend(value: str) -> T2VBackend: - """Resolve a CLI/UI backend key.""" - try: - return BACKENDS[value] - except KeyError as exc: - raise ValueError( - f"Unknown backend {value!r}. Available backends: {', '.join(BACKENDS)}." - ) from exc - - -def backend_choices() -> tuple[str, ...]: - """Return stable CLI choices.""" - return tuple(BACKENDS) - - -def backend_metadata() -> list[dict[str, Any]]: - """Return browser-safe backend names and app-owned presets.""" - return [ - { - "key": backend.key, - "label": backend.label, - "default_preset": backend.default_preset_name, - "presets": backend.preset_names, - } - for backend in BACKENDS.values() - ] diff --git a/apps/t2v_demo/presets.py b/apps/t2v_demo/presets.py deleted file mode 100644 index 592a3bd73..000000000 --- a/apps/t2v_demo/presets.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""App-owned T2V pipeline presets, independent of workspace integrations.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any - -from torch import Tensor - -from flashdreams.infra.config import derive_config -from flashdreams.infra.diffusion.model import DiffusionModelConfig -from flashdreams.infra.diffusion.scheduler import FlowMatchUniPCSchedulerConfig -from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig -from flashdreams.recipes.cosmos.pipeline import CosmosInferencePipelineConfig -from flashdreams.recipes.cosmos.transformer import CosmosTransformerConfig -from flashdreams.recipes.cosmos.transformer.impl.network import ( - CosmosDiTNetworkConfig, -) -from flashdreams.recipes.cosmos.transformer.impl.network import ( - state_dict_transform as cosmos_state_dict_transform, -) -from flashdreams.recipes.taehv import TeahvVAEDecoderConfig -from flashdreams.recipes.wan import ( - Wan21TransformerConfig, - WanDiTNetwork1pt3BConfig, - WanInferencePipelineConfig, - WanVAEDecoderConfig, -) - - -@dataclass(frozen=True, slots=True) -class T2VPreset: - """One app-supported text-to-video pipeline configuration.""" - - name: str - pipeline: Any - prompt: str - total_blocks: int - pixel_height: int - pixel_width: int - fps: int - - -CAUSAL_FORCING_PROMPT = "A cinematic closeup of a reindeer in a snowy forest at sunset." -SELF_FORCING_PROMPT = "A stylish woman strolls down a neon-lit Tokyo street at night." -COSMOS_PROMPT = ( - "A high-definition video of a robotic arm welding in an industrial workshop." -) - - -def _wan_state_dict_transform(state_dict: dict[str, Any]) -> dict[str, Tensor]: - """Normalize upstream Causal/Self-Forcing checkpoint wrapper keys.""" - state_dict = state_dict.get( - "generator_ema", state_dict.get("generator", state_dict) - ) - out: dict[str, Tensor] = {} - for key, value in state_dict.items(): - key = key.removeprefix("model.").removeprefix("net.") - out[key.removeprefix("_fsdp_wrapped_module.")] = value - return out - - -CAUSAL_FORCING_PIPELINE = WanInferencePipelineConfig( - name="causal-forcing-wan2.1-t2v-1.3b-chunkwise", - enable_sync_and_profile=True, - encoder=None, - decoder=WanVAEDecoderConfig(), - diffusion_model=DiffusionModelConfig( - seed=42, - transformer=Wan21TransformerConfig( - network=WanDiTNetwork1pt3BConfig( - patch_embedding_type="conv3d", cp_method="ring" - ), - checkpoint_path="https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/chunkwise/causal_forcing.pt", - state_dict_transform=_wan_state_dict_transform, - batch_shape=(), - len_t=3, - guidance_scale=1.0, - window_size_t=21, - sink_size_t=0, - stamp_image_latent=False, - compile_network=True, - ), - scheduler=FlowMatchSchedulerConfig( - num_inference_steps=4, - denoising_timesteps=[1000, 750, 500, 250], - warp_denoising_step=True, - shift=5.0, - sigma_min=0.0, - extra_one_step=True, - num_train_timesteps=1000, - ), - ), -) -CAUSAL_FORCING_FRAMEWISE_PIPELINE = derive_config( - CAUSAL_FORCING_PIPELINE, - name="causal-forcing-wan2.1-t2v-1.3b-framewise", - diffusion_model=dict( - transformer=dict( - checkpoint_path="https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/framewise/causal_forcing.pt", - len_t=1, - ), - ), -) - -SELF_FORCING_PIPELINE = WanInferencePipelineConfig( - name="self-forcing-wan2.1-t2v-1.3b", - enable_sync_and_profile=True, - encoder=None, - decoder=WanVAEDecoderConfig(), - diffusion_model=DiffusionModelConfig( - seed=42, - transformer=Wan21TransformerConfig( - network=WanDiTNetwork1pt3BConfig( - patch_embedding_type="conv3d", cp_method="ring" - ), - checkpoint_path="https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt", - state_dict_transform=_wan_state_dict_transform, - batch_shape=(), - len_t=3, - guidance_scale=1.0, - window_size_t=21, - sink_size_t=0, - stamp_image_latent=False, - compile_network=True, - ), - scheduler=FlowMatchSchedulerConfig( - num_inference_steps=4, - denoising_timesteps=[1000, 750, 500, 250], - warp_denoising_step=True, - shift=8.0, - sigma_min=0.0, - extra_one_step=True, - num_train_timesteps=1000, - ), - ), -) -SELF_FORCING_TAEHV_PIPELINE = derive_config( - SELF_FORCING_PIPELINE, - name="self-forcing-wan2.1-t2v-1.3b-taehv", - decoder=TeahvVAEDecoderConfig(), -) -SELF_FORCING_REROPE_PIPELINE = derive_config( - SELF_FORCING_PIPELINE, - name="self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope", - diffusion_model=dict( - seed=0, - transformer=dict( - window_size_t=7, - sink_size_t=5, - compile_network=False, - use_cuda_graph=False, - network=dict(apply_rope_before_kvcache=False), - ), - ), -) - -COSMOS_PIPELINE = CosmosInferencePipelineConfig( - name="cosmos2-t2v-2b-720p", - enable_sync_and_profile=True, - encoder=None, - decoder=WanVAEDecoderConfig(), - diffusion_model=DiffusionModelConfig( - seed=42, - transformer=CosmosTransformerConfig( - network=CosmosDiTNetworkConfig(cp_method="ring"), - checkpoint_path="https://huggingface.co/nvidia/Cosmos-Predict2.5-2B/blob/main/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt", - state_dict_transform=cosmos_state_dict_transform, - batch_shape=(), - len_t=24, - window_size_t=24, - guidance_scale=8.0, - compile_network=True, - use_cuda_graph=False, - ), - scheduler=FlowMatchUniPCSchedulerConfig( - num_inference_steps=35, shift=5.0, use_kerras_sigma=True, enable_tqdm=True - ), - ), -) - -PRESETS: dict[str, T2VPreset] = { - preset.name: preset - for preset in ( - T2VPreset( - name=CAUSAL_FORCING_PIPELINE.name, - pipeline=CAUSAL_FORCING_PIPELINE, - prompt=CAUSAL_FORCING_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, - ), - T2VPreset( - name=CAUSAL_FORCING_FRAMEWISE_PIPELINE.name, - pipeline=CAUSAL_FORCING_FRAMEWISE_PIPELINE, - prompt=CAUSAL_FORCING_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, - ), - T2VPreset( - name=SELF_FORCING_PIPELINE.name, - pipeline=SELF_FORCING_PIPELINE, - prompt=SELF_FORCING_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, - ), - T2VPreset( - name=SELF_FORCING_TAEHV_PIPELINE.name, - pipeline=SELF_FORCING_TAEHV_PIPELINE, - prompt=SELF_FORCING_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, - ), - T2VPreset( - name=SELF_FORCING_REROPE_PIPELINE.name, - pipeline=SELF_FORCING_REROPE_PIPELINE, - prompt=SELF_FORCING_PROMPT, - total_blocks=80, - pixel_height=480, - pixel_width=832, - fps=16, - ), - T2VPreset( - name=COSMOS_PIPELINE.name, - pipeline=COSMOS_PIPELINE, - prompt=COSMOS_PROMPT, - total_blocks=1, - pixel_height=720, - pixel_width=1280, - fps=16, - ), - ) -} diff --git a/apps/t2v_demo/pyproject.toml b/apps/t2v_demo/pyproject.toml index e357fd44d..64453965c 100644 --- a/apps/t2v_demo/pyproject.toml +++ b/apps/t2v_demo/pyproject.toml @@ -10,7 +10,13 @@ name = "flashdreams-t2v-demo" version = "0.1.0" description = "FlashDreams text-to-video runtime demo launcher" requires-python = ">=3.10" -dependencies = ["flashdreams[serving]", "flashdreams-t2v"] +dependencies = [ + "flashdreams[serving]", + "flashdreams-causal-forcing", + "flashdreams-cosmos-predict2", + "flashdreams-self-forcing", + "flashdreams-t2v", +] [project.entry-points."flashdreams.runner_configs"] t2v = "t2v_demo.runner:RUNNER_T2V" @@ -20,6 +26,9 @@ t2v = "t2v_demo.app:create_app" [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-causal-forcing = { workspace = true } +flashdreams-cosmos-predict2 = { workspace = true } +flashdreams-self-forcing = { workspace = true } flashdreams-t2v = { workspace = true } [tool.setuptools] diff --git a/apps/t2v_demo/runner.py b/apps/t2v_demo/runner.py index 8fe620568..fd539f9fb 100644 --- a/apps/t2v_demo/runner.py +++ b/apps/t2v_demo/runner.py @@ -13,7 +13,7 @@ from flashdreams.infra.runner import Runner, RunnerConfig -from .backends import backend_choices, resolve_backend +from .runtime import backend_choices, default_pipeline @dataclass(kw_only=True) @@ -25,9 +25,7 @@ class T2VDemoRunnerConfig(RunnerConfig): "t2v_demo.launch:LAUNCH_CAPABILITY" ) pipeline: Annotated[Any, tyro.conf.Suppress] = field( - default_factory=lambda: resolve_backend("causal-forcing") - .resolve_runner() - .pipeline + default_factory=default_pipeline ) backend: str = "causal-forcing" """Backend key: one of ``causal-forcing``, ``cosmos-predict2``, or ``self-forcing``.""" diff --git a/apps/t2v_demo/runtime.py b/apps/t2v_demo/runtime.py index 8e149ceba..8f0040d15 100644 --- a/apps/t2v_demo/runtime.py +++ b/apps/t2v_demo/runtime.py @@ -1,11 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Compatibility bridge from the legacy T2V demo registry to the shared shell.""" +"""Compatibility bridge from ``flashdreams-run t2v`` to integration apps.""" from __future__ import annotations -from t2v.t2v import ( +from dataclasses import dataclass, replace +from importlib import import_module +from importlib.metadata import entry_points +from typing import Any + +from t2v import ( FIELD_FPS, FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH, @@ -17,36 +22,181 @@ T2VRuntime, T2VScenario, T2VSession, + model_config_from_runner, ) -from .backends import T2VBackend, resolve_backend +from flashdreams.demo import Application, DemoAdapterApplication + + +@dataclass(frozen=True, slots=True) +class T2VBackendBridge: + """One legacy backend key routed to one integration-owned app.""" + + key: str + label: str + app_slug: str + app_module: str + config_module: str + + +_BACKENDS: dict[str, T2VBackendBridge] = { + "causal-forcing": T2VBackendBridge( + key="causal-forcing", + label="Causal-Forcing (Wan 2.1)", + app_slug="causal-forcing-t2v", + app_module="causal_forcing.t2v.app", + config_module="causal_forcing.config", + ), + "cosmos-predict2": T2VBackendBridge( + key="cosmos-predict2", + label="Cosmos Predict2", + app_slug="cosmos-predict2-t2v", + app_module="cosmos_predict2.t2v.app", + config_module="cosmos_predict2.config", + ), + "self-forcing": T2VBackendBridge( + key="self-forcing", + label="Self-Forcing (Wan 2.1)", + app_slug="self-forcing-t2v", + app_module="self_forcing.t2v.app", + config_module="self_forcing.config", + ), +} + + +def backend_choices() -> tuple[str, ...]: + """Return stable legacy CLI backend choices.""" + return tuple(_BACKENDS) + + +def backend_metadata() -> list[dict[str, Any]]: + """Return browser-safe backend metadata derived from integration configs.""" + metadata: list[dict[str, Any]] = [] + for bridge in _BACKENDS.values(): + model = model_from_backend(bridge.key) + metadata.append( + { + "key": bridge.key, + "label": bridge.label, + "default_preset": model.preset_id, + "presets": _t2v_preset_ids(bridge), + "application": bridge.app_slug, + } + ) + return metadata + + +def default_pipeline() -> Any: + """Return the legacy runner's suppressed default pipeline value.""" + return model_from_backend("causal-forcing").pipeline def model_from_backend( - backend: str | T2VBackend, + backend: str, preset_id: str | None = None, ) -> T2VModelConfig: - """Adapt the existing backend registry to one neutral T2V model config.""" - resolved = resolve_backend(backend) if isinstance(backend, str) else backend - preset = resolved.resolve_runner(preset_id) - return T2VModelConfig( - model_id="flashdreams-t2v", - preset_id=preset.name, - pipeline=preset.pipeline, - prompt=preset.prompt, - total_blocks=preset.total_blocks, - pixel_height=preset.pixel_height, - pixel_width=preset.pixel_width, - fps=preset.fps, - runtime_options={"backend": resolved.key}, + """Resolve a legacy backend/preset selector to an integration-owned model.""" + bridge = _resolve_backend(backend) + if preset_id is None: + return _with_legacy_backend_option( + _default_model_from_application(bridge), bridge + ) + return _with_legacy_backend_option( + _model_from_runner_config(bridge, preset_id), bridge ) def make_adapter(backend: str, preset_id: str | None = None) -> T2VDemoAdapter: - """Build an adapter from a CLI/UI backend key.""" + """Build an adapter from a legacy CLI/UI backend key.""" return T2VDemoAdapter(model=model_from_backend(backend, preset_id)) +def _resolve_backend(value: str) -> T2VBackendBridge: + try: + return _BACKENDS[value] + except KeyError as exc: + raise ValueError( + f"Unknown backend {value!r}. Available backends: {', '.join(_BACKENDS)}." + ) from exc + + +def _default_model_from_application(bridge: T2VBackendBridge) -> T2VModelConfig: + app = _load_application(bridge) + if not isinstance(app, DemoAdapterApplication): + raise TypeError( + f"T2V application {bridge.app_slug!r} must return " + f"DemoAdapterApplication, got {type(app).__name__}." + ) + adapter = app.adapter + model = getattr(adapter, "model", None) + if not isinstance(model, T2VModelConfig): + raise TypeError(f"T2V application {bridge.app_slug!r} must use T2VDemoAdapter.") + return model + + +def _load_application(bridge: T2VBackendBridge) -> Application: + for entry_point in entry_points(group="flashdreams.applications"): + if entry_point.name == bridge.app_slug: + factory = entry_point.load() + return _coerce_application( + factory() if callable(factory) else factory, + app_slug=bridge.app_slug, + ) + factory = getattr(import_module(bridge.app_module), "create_app") + return _coerce_application(factory(), app_slug=bridge.app_slug) + + +def _coerce_application(value: object, *, app_slug: str) -> Application: + if not isinstance(value, Application): + raise TypeError( + f"T2V application {app_slug!r} must return Application, " + f"got {type(value).__name__}." + ) + return value + + +def _model_from_runner_config( + bridge: T2VBackendBridge, preset_id: str +) -> T2VModelConfig: + runner_configs = getattr(import_module(bridge.config_module), "RUNNER_CONFIGS") + try: + runner = runner_configs[preset_id] + except KeyError as exc: + raise ValueError( + f"Unknown {bridge.key} preset {preset_id!r}. Available presets: " + f"{', '.join(_t2v_preset_ids(bridge))}." + ) from exc + if not _is_t2v_preset(runner.runner_name): + raise ValueError( + f"Preset {preset_id!r} is not a T2V preset for backend {bridge.key!r}." + ) + default_model = _default_model_from_application(bridge) + return model_config_from_runner(model_id=default_model.model_id, runner=runner) + + +def _t2v_preset_ids(bridge: T2VBackendBridge) -> tuple[str, ...]: + runner_configs = getattr(import_module(bridge.config_module), "RUNNER_CONFIGS") + return tuple(name for name in runner_configs if _is_t2v_preset(name)) + + +def _is_t2v_preset(name: str) -> bool: + return "-t2v-" in name + + +def _with_legacy_backend_option( + model: T2VModelConfig, + bridge: T2VBackendBridge, +) -> T2VModelConfig: + return replace( + model, + runtime_options={ + **model.runtime_options, + "backend": bridge.key, + "application": bridge.app_slug, + }, + ) + + __all__ = [ "FIELD_FPS", "FIELD_PIXEL_HEIGHT", @@ -54,11 +204,15 @@ def make_adapter(backend: str, preset_id: str | None = None) -> T2VDemoAdapter: "FIELD_PROMPT", "FIELD_TOTAL_BLOCKS", "T2VDemoAdapter", + "T2VBackendBridge", "T2VInputProvider", "T2VModelConfig", "T2VRuntime", "T2VScenario", "T2VSession", + "backend_choices", + "backend_metadata", + "default_pipeline", "make_adapter", "model_from_backend", ] diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py index fc4d928c6..1cf6da22c 100644 --- a/apps/t2v_demo/tests/test_runner.py +++ b/apps/t2v_demo/tests/test_runner.py @@ -10,7 +10,7 @@ import tomli from t2v_demo import app from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig -from t2v_demo.runtime import model_from_backend +from t2v_demo.runtime import backend_metadata, model_from_backend from flashdreams.demo import ( Application, @@ -37,6 +37,13 @@ def test_t2v_registers_application_entry_point() -> None: } +def test_t2v_demo_no_longer_owns_backend_presets() -> None: + package_dir = Path(__file__).parents[1] + + assert not (package_dir / "backends.py").exists() + assert not (package_dir / "presets.py").exists() + + def test_t2v_create_app_exposes_public_application() -> None: public_app = app.create_app( T2VDemoRunnerConfig( @@ -52,9 +59,10 @@ def test_t2v_create_app_exposes_public_application() -> None: assert isinstance(public_app, Application) assert isinstance(public_app, DemoAdapterApplication) spec = public_app.spec - assert spec.model_id == "flashdreams-t2v" + assert spec.model_id == "self-forcing-t2v" assert spec.config is not None assert spec.config.runtime_options["backend"] == "self-forcing" + assert spec.config.runtime_options["application"] == "self-forcing-t2v" scenario = spec.scenario assert isinstance(scenario, dict) assert scenario["prompt"] == "A waterfall" @@ -65,10 +73,31 @@ def test_t2v_backend_bridge_builds_neutral_model_config() -> None: model = model_from_backend("self-forcing") assert isinstance(model, t2v_shell.T2VModelConfig) - assert model.model_id == "flashdreams-t2v" + assert model.model_id == "self-forcing-t2v" assert model.runtime_options["backend"] == "self-forcing" +def test_t2v_backend_bridge_supports_integration_owned_preset() -> None: + model = model_from_backend( + "self-forcing", + "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope", + ) + + assert model.model_id == "self-forcing-t2v" + assert model.preset_id == "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope" + assert model.total_blocks == 80 + + +def test_t2v_backend_metadata_is_derived_from_integrations() -> None: + metadata = {item["key"]: item for item in backend_metadata()} + + assert metadata["self-forcing"]["default_preset"] == ( + "self-forcing-wan2.1-t2v-1.3b" + ) + assert "self-forcing-wan2.1-t2v-1.3b-taehv" in metadata["self-forcing"]["presets"] + assert metadata["cosmos-predict2"]["application"] == "cosmos-predict2-t2v" + + def test_runner_mp4_launch_uses_demo_entrypoint( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/integrations/causal_forcing/causal_forcing/t2v/__init__.py b/integrations/causal_forcing/causal_forcing/t2v/__init__.py index 175fd535e..0507d5c35 100644 --- a/integrations/causal_forcing/causal_forcing/t2v/__init__.py +++ b/integrations/causal_forcing/causal_forcing/t2v/__init__.py @@ -3,6 +3,6 @@ """Causal-Forcing T2V public demo app.""" -from causal_forcing.t2v.app import MODEL, createApp, create_app +from causal_forcing.t2v.app import MODEL, create_app, createApp __all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/causal_forcing/causal_forcing/t2v/app.py b/integrations/causal_forcing/causal_forcing/t2v/app.py index ab71a6d31..631e4cc0e 100644 --- a/integrations/causal_forcing/causal_forcing/t2v/app.py +++ b/integrations/causal_forcing/causal_forcing/t2v/app.py @@ -3,21 +3,14 @@ """Public T2V app entry for the default Causal-Forcing model.""" -from flashdreams.demo import Application -from t2v import T2VModelConfig, create_t2v_application +from t2v import create_t2v_application, model_config_from_runner -from causal_forcing.config import PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE -from causal_forcing.runner import DEFAULT_T2V_PROMPT +from causal_forcing.config import RUNNER_WAN21_T2V_1PT3B_CHUNKWISE +from flashdreams.demo import Application -MODEL = T2VModelConfig( +MODEL = model_config_from_runner( model_id="causal-forcing-t2v", - preset_id=PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE.name, - pipeline=PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE, - prompt=DEFAULT_T2V_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, + runner=RUNNER_WAN21_T2V_1PT3B_CHUNKWISE, ) diff --git a/integrations/causal_forcing/tests/test_smoke.py b/integrations/causal_forcing/tests/test_smoke.py index 39c5ff9d4..644d23d78 100644 --- a/integrations/causal_forcing/tests/test_smoke.py +++ b/integrations/causal_forcing/tests/test_smoke.py @@ -32,7 +32,7 @@ import pytest import tomli as tomllib from causal_forcing import config as config_mod -from causal_forcing.config import PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE, RUNNER_CONFIGS +from causal_forcing.config import RUNNER_CONFIGS, RUNNER_WAN21_T2V_1PT3B_CHUNKWISE from causal_forcing.t2v.app import MODEL, create_app, createApp from flashdreams.demo import Application, DemoAdapterApplication @@ -108,8 +108,8 @@ def test_t2v_app_uses_default_pipeline_config() -> None: assert isinstance(public_app, Application) assert isinstance(public_app, DemoAdapterApplication) assert MODEL.model_id == "causal-forcing-t2v" - assert MODEL.preset_id == PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE.name - assert MODEL.pipeline is PIPELINE_WAN21_T2V_1PT3B_CHUNKWISE + assert MODEL.preset_id == RUNNER_WAN21_T2V_1PT3B_CHUNKWISE.runner_name + assert MODEL.pipeline is RUNNER_WAN21_T2V_1PT3B_CHUNKWISE.pipeline assert public_app.spec.model_id == MODEL.model_id assert public_app.spec.preset_id == MODEL.preset_id diff --git a/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py b/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py index a72ebbf6e..68fcfb392 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py +++ b/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py @@ -3,6 +3,6 @@ """Cosmos Predict2 T2V public demo app.""" -from cosmos_predict2.t2v.app import MODEL, createApp, create_app +from cosmos_predict2.t2v.app import MODEL, create_app, createApp __all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py b/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py index a9467fc4b..8e2716ff7 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py +++ b/integrations/cosmos_predict2/cosmos_predict2/t2v/app.py @@ -3,21 +3,14 @@ """Public T2V app entry for the default Cosmos Predict2 model.""" -from flashdreams.demo import Application -from t2v import T2VModelConfig, create_t2v_application +from t2v import create_t2v_application, model_config_from_runner -from cosmos_predict2.config import PIPELINE_COSMOS2_T2V_2B_720P -from cosmos_predict2.runner import DEFAULT_PROMPT +from cosmos_predict2.config import RUNNER_COSMOS2_T2V_2B_720P +from flashdreams.demo import Application -MODEL = T2VModelConfig( +MODEL = model_config_from_runner( model_id="cosmos-predict2-t2v", - preset_id=PIPELINE_COSMOS2_T2V_2B_720P.name, - pipeline=PIPELINE_COSMOS2_T2V_2B_720P, - prompt=DEFAULT_PROMPT, - total_blocks=1, - pixel_height=720, - pixel_width=1280, - fps=16, + runner=RUNNER_COSMOS2_T2V_2B_720P, ) diff --git a/integrations/cosmos_predict2/tests/test_t2v_app.py b/integrations/cosmos_predict2/tests/test_t2v_app.py index fa3f3483e..1a50b47ed 100644 --- a/integrations/cosmos_predict2/tests/test_t2v_app.py +++ b/integrations/cosmos_predict2/tests/test_t2v_app.py @@ -7,8 +7,8 @@ import pytest import tomli as tomllib -from cosmos_predict2.config import PIPELINE_COSMOS2_T2V_2B_720P -from cosmos_predict2.t2v.app import MODEL, createApp, create_app +from cosmos_predict2.config import RUNNER_COSMOS2_T2V_2B_720P +from cosmos_predict2.t2v.app import MODEL, create_app, createApp from flashdreams.demo import Application, DemoAdapterApplication @@ -25,8 +25,8 @@ def test_t2v_app_uses_default_pipeline_config() -> None: assert isinstance(public_app, Application) assert isinstance(public_app, DemoAdapterApplication) assert MODEL.model_id == "cosmos-predict2-t2v" - assert MODEL.preset_id == PIPELINE_COSMOS2_T2V_2B_720P.name - assert MODEL.pipeline is PIPELINE_COSMOS2_T2V_2B_720P + assert MODEL.preset_id == RUNNER_COSMOS2_T2V_2B_720P.runner_name + assert MODEL.pipeline is RUNNER_COSMOS2_T2V_2B_720P.pipeline assert public_app.spec.model_id == MODEL.model_id assert public_app.spec.preset_id == MODEL.preset_id @@ -38,6 +38,4 @@ def test_application_entry_point_matches_module_literal() -> None: meta = tomllib.load(fh) entries = meta["project"]["entry-points"][APPLICATION_ENTRY_POINT_GROUP] - assert entries == { - "cosmos-predict2-t2v": "cosmos_predict2.t2v.app:create_app" - } + assert entries == {"cosmos-predict2-t2v": "cosmos_predict2.t2v.app:create_app"} diff --git a/integrations/self_forcing/self_forcing/t2v/__init__.py b/integrations/self_forcing/self_forcing/t2v/__init__.py index 1ed51e404..57b5c47fb 100644 --- a/integrations/self_forcing/self_forcing/t2v/__init__.py +++ b/integrations/self_forcing/self_forcing/t2v/__init__.py @@ -3,6 +3,6 @@ """Self-Forcing T2V public demo app.""" -from self_forcing.t2v.app import MODEL, createApp, create_app +from self_forcing.t2v.app import MODEL, create_app, createApp __all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/self_forcing/self_forcing/t2v/app.py b/integrations/self_forcing/self_forcing/t2v/app.py index 2eb64cf35..bbb10e54c 100644 --- a/integrations/self_forcing/self_forcing/t2v/app.py +++ b/integrations/self_forcing/self_forcing/t2v/app.py @@ -3,21 +3,14 @@ """Public T2V app entry for the default Self-Forcing model.""" -from flashdreams.demo import Application -from t2v import T2VModelConfig, create_t2v_application +from t2v import create_t2v_application, model_config_from_runner -from self_forcing.config import PIPELINE_WAN21_T2V_1PT3B -from self_forcing.runner import DEFAULT_T2V_PROMPT +from flashdreams.demo import Application +from self_forcing.config import RUNNER_WAN21_T2V_1PT3B -MODEL = T2VModelConfig( +MODEL = model_config_from_runner( model_id="self-forcing-t2v", - preset_id=PIPELINE_WAN21_T2V_1PT3B.name, - pipeline=PIPELINE_WAN21_T2V_1PT3B, - prompt=DEFAULT_T2V_PROMPT, - total_blocks=60, - pixel_height=480, - pixel_width=832, - fps=16, + runner=RUNNER_WAN21_T2V_1PT3B, ) diff --git a/integrations/self_forcing/tests/test_smoke.py b/integrations/self_forcing/tests/test_smoke.py index 3e6377fe9..b5c482d8c 100644 --- a/integrations/self_forcing/tests/test_smoke.py +++ b/integrations/self_forcing/tests/test_smoke.py @@ -24,7 +24,7 @@ import pytest import tomli as tomllib from self_forcing import config as config_mod -from self_forcing.config import PIPELINE_WAN21_T2V_1PT3B, RUNNER_CONFIGS +from self_forcing.config import RUNNER_CONFIGS, RUNNER_WAN21_T2V_1PT3B from self_forcing.t2v.app import MODEL, create_app, createApp from flashdreams.demo import Application, DemoAdapterApplication @@ -100,8 +100,8 @@ def test_t2v_app_uses_default_pipeline_config() -> None: assert isinstance(public_app, Application) assert isinstance(public_app, DemoAdapterApplication) assert MODEL.model_id == "self-forcing-t2v" - assert MODEL.preset_id == PIPELINE_WAN21_T2V_1PT3B.name - assert MODEL.pipeline is PIPELINE_WAN21_T2V_1PT3B + assert MODEL.preset_id == RUNNER_WAN21_T2V_1PT3B.runner_name + assert MODEL.pipeline is RUNNER_WAN21_T2V_1PT3B.pipeline assert public_app.spec.model_id == MODEL.model_id assert public_app.spec.preset_id == MODEL.preset_id diff --git a/uv.lock b/uv.lock index f8190bfd3..c37b144be 100644 --- a/uv.lock +++ b/uv.lock @@ -1442,12 +1442,18 @@ version = "0.1.0" source = { editable = "apps/t2v_demo" } dependencies = [ { name = "flashdreams", extra = ["serving"] }, + { name = "flashdreams-causal-forcing" }, + { name = "flashdreams-cosmos-predict2" }, + { name = "flashdreams-self-forcing" }, { name = "flashdreams-t2v" }, ] [package.metadata] requires-dist = [ { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, + { name = "flashdreams-causal-forcing", editable = "integrations/causal_forcing" }, + { name = "flashdreams-cosmos-predict2", editable = "integrations/cosmos_predict2" }, + { name = "flashdreams-self-forcing", editable = "integrations/self_forcing" }, { name = "flashdreams-t2v", editable = "apps/t2v" }, ] From aa69c9b4df7c13e5139d1694c966e3ea8dc7bdac Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 11:03:02 +0000 Subject: [PATCH 22/33] Fix public runner cleanup edge cases Reject closed external hosts before app initialization, avoid direct runner-thread application cleanup when worker dispatch is unavailable, and preserve primary async run failures when cancellation arrives during cleanup. --- apps/t2v/t2v.py | 235 +++++++++++------- apps/t2v/tests/test_t2v_shell.py | 36 ++- apps/t2v_demo/app.py | 6 +- apps/t2v_demo/runtime.py | 12 +- flashdreams/flashdreams/demo/runner.py | 81 +++++- .../tests/test_demo_application_api.py | 45 +++- 6 files changed, 303 insertions(+), 112 deletions(-) diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index 583ea8b15..464d42480 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -31,6 +31,7 @@ InputField, ModelAdapter, StepRequest, + StepRequirements, ) from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.demo import ( @@ -57,6 +58,8 @@ FIELD_PIXEL_WIDTH = "pixel_width" FIELD_FPS = "fps" +_DOWNLOAD_ARTIFACT_DIR = Path("outputs/t2v-webrtc") + @dataclass(frozen=True, kw_only=True, slots=True) class T2VModelConfig: @@ -77,16 +80,23 @@ def __post_init__(self) -> None: raise ValueError("T2VModelConfig.model_id must be non-empty.") if self.preset_id is not None and not self.preset_id.strip(): raise ValueError("T2VModelConfig.preset_id must be non-empty when set.") - if not self.prompt.strip(): - raise ValueError("T2VModelConfig.prompt must be non-empty.") - _validate_positive_int(self.total_blocks, name=FIELD_TOTAL_BLOCKS) - _validate_positive_int(self.pixel_height, name=FIELD_PIXEL_HEIGHT) - _validate_positive_int(self.pixel_width, name=FIELD_PIXEL_WIDTH) - _validate_positive_int(self.fps, name=FIELD_FPS) + # T2VScenario owns prompt and geometry validation so that model + # defaults and launch-time overrides are held to the same rules. + self.default_scenario() object.__setattr__( self, "runtime_options", freeze_mapping(self.runtime_options) ) + def default_scenario(self) -> T2VScenario: + """Return this model's scenario before any launch-time override.""" + return T2VScenario( + prompt=self.prompt, + total_blocks=self.total_blocks, + pixel_height=self.pixel_height, + pixel_width=self.pixel_width, + fps=self.fps, + ) + @dataclass(frozen=True, kw_only=True, slots=True) class T2VRunDefaults: @@ -114,6 +124,17 @@ def __post_init__(self) -> None: self, "runtime_options", freeze_mapping(self.runtime_options) ) + def scenario_overrides(self) -> dict[str, object]: + """Return only the scenario fields this launch explicitly overrides.""" + overrides = { + FIELD_PROMPT: self.prompt, + FIELD_TOTAL_BLOCKS: self.total_blocks, + FIELD_PIXEL_HEIGHT: self.pixel_height, + FIELD_PIXEL_WIDTH: self.pixel_width, + FIELD_FPS: self.fps, + } + return {name: value for name, value in overrides.items() if value is not None} + @dataclass(frozen=True, kw_only=True, slots=True) class T2VScenario: @@ -125,6 +146,42 @@ class T2VScenario: pixel_width: int fps: int + def __post_init__(self) -> None: + if not self.prompt.strip(): + raise ValueError("A non-empty text-to-video prompt is required.") + _validate_positive_int(self.total_blocks, name=FIELD_TOTAL_BLOCKS) + _validate_positive_int(self.pixel_height, name=FIELD_PIXEL_HEIGHT) + _validate_positive_int(self.pixel_width, name=FIELD_PIXEL_WIDTH) + _validate_positive_int(self.fps, name=FIELD_FPS) + + @classmethod + def from_mapping( + cls, + source: Mapping[str, object], + *, + defaults: T2VScenario | None = None, + ) -> T2VScenario: + """Build a scenario from runtime values, falling back to ``defaults``.""" + merged: dict[str, object] = {} if defaults is None else defaults.to_mapping() + merged.update(source) + return cls( + prompt=str(merged[FIELD_PROMPT]).strip(), + total_blocks=_int_value(merged[FIELD_TOTAL_BLOCKS]), + pixel_height=_int_value(merged[FIELD_PIXEL_HEIGHT]), + pixel_width=_int_value(merged[FIELD_PIXEL_WIDTH]), + fps=_int_value(merged[FIELD_FPS]), + ) + + def to_mapping(self) -> dict[str, object]: + """Return the scenario as the runtime's global conditioning values.""" + return { + FIELD_PROMPT: self.prompt, + FIELD_TOTAL_BLOCKS: self.total_blocks, + FIELD_PIXEL_HEIGHT: self.pixel_height, + FIELD_PIXEL_WIDTH: self.pixel_width, + FIELD_FPS: self.fps, + } + class T2VDemoAdapter(ModelAdapter): """Model adapter shared by replay and WebRTC T2V launch paths.""" @@ -135,8 +192,11 @@ class T2VDemoAdapter(ModelAdapter): ) canonical_input_schema = CanonicalInputSchema() - def __init__(self, *, model: T2VModelConfig) -> None: + def __init__( + self, *, model: T2VModelConfig, write_download_artifact: bool = False + ) -> None: self.model = model + self.write_download_artifact = write_download_artifact @property def model_id(self) -> str: @@ -172,22 +232,21 @@ def validate_config(self, config: InferenceConfig) -> None: ) def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: - scenario = _scenario_from_value(spec.scenario, self.model) + source = spec.scenario if isinstance(spec.scenario, Mapping) else {} + scenario = T2VScenario.from_mapping( + source, defaults=self.model.default_scenario() + ) return PreparedScenario( - initial_inputs=InferenceInput( - global_conditioning={ - FIELD_PROMPT: scenario.prompt, - FIELD_TOTAL_BLOCKS: scenario.total_blocks, - FIELD_PIXEL_HEIGHT: scenario.pixel_height, - FIELD_PIXEL_WIDTH: scenario.pixel_width, - FIELD_FPS: scenario.fps, - } - ) + initial_inputs=InferenceInput(global_conditioning=scenario.to_mapping()) ) def create_runtime(self, config: InferenceConfig) -> "T2VRuntime": self.validate_config(config) - return T2VRuntime(config=config, model=self.model) + return T2VRuntime( + config=config, + model=self.model, + write_download_artifact=self.write_download_artifact, + ) def create_model_input_provider( self, spec: DemoSpec, scenario: PreparedScenario @@ -213,7 +272,7 @@ def prepare_initial_input(self) -> InferenceInput: return self._initial_inputs def prepare_step( - self, *, request: Any, user_window: UserInputWindow + self, *, request: StepRequirements, user_window: UserInputWindow ) -> PreparedStep: del request, user_window return PreparedStep(inference_input=InferenceInput()) @@ -229,9 +288,16 @@ def close(self) -> None: class T2VRuntime: """One heavyweight selected pipeline, reusable across demo sessions.""" - def __init__(self, *, config: InferenceConfig, model: T2VModelConfig) -> None: + def __init__( + self, + *, + config: InferenceConfig, + model: T2VModelConfig, + write_download_artifact: bool = False, + ) -> None: self.config = config self.model = model + self._write_download_artifact = write_download_artifact pipeline_config = model.pipeline if config.compile is not None: from flashdreams.infra.config import derive_config @@ -262,7 +328,12 @@ def latest_artifact(self) -> tuple[Path, T2VScenario] | None: def start_session(self, inputs: InferenceInput) -> "T2VSession": return T2VSession( - pipeline=self.pipeline, scenario=_scenario_from_inputs(inputs), runtime=self + pipeline=self.pipeline, + scenario=T2VScenario.from_mapping(inputs.global_conditioning), + runtime=self, + artifact_dir=( + _DOWNLOAD_ARTIFACT_DIR if self._write_download_artifact else None + ), ) def close(self) -> None: @@ -277,25 +348,40 @@ class T2VSession(InferenceSession): """A cache-isolated T2V session that yields chunks as they are generated.""" def __init__( - self, *, pipeline: Any, scenario: T2VScenario, runtime: T2VRuntime + self, + *, + pipeline: Any, + scenario: T2VScenario, + runtime: T2VRuntime, + artifact_dir: Path | None = None, ) -> None: self.pipeline = pipeline self.scenario = scenario self._runtime = runtime - self._artifact_path = Path("outputs/t2v-webrtc") / f"{uuid4()}.mp4" - self._artifact_path.parent.mkdir(parents=True, exist_ok=True) - # Transitional support for the existing WebRTC download endpoint. - # Replay modes still deliver primary output through the shared OutputSink. - self._artifact_output = Mp4VideoOutputTarget( - output_path=self._artifact_path, fps=scenario.fps, output_layout="tchw" - ) - self._artifact_output.open() + # Transitional second encode that backs the WebRTC download endpoint. + # Replay modes deliver primary output through the shared OutputSink and + # leave this disabled so they do not encode every chunk twice. + self._artifact_path: Path | None = None + self._artifact_output: Mp4VideoOutputTarget | None = None + if artifact_dir is not None: + artifact_dir.mkdir(parents=True, exist_ok=True) + self._artifact_path = artifact_dir / f"{uuid4()}.mp4" + self._artifact_output = Mp4VideoOutputTarget( + output_path=self._artifact_path, + fps=scenario.fps, + output_layout="tchw", + ) + self._artifact_output.open() self._step_index = 0 self._closed = False self._output_stream = VideoOutputStream( postprocess_stream=None, output_layout="tchw" ) - assert isinstance(pipeline.decoder, StreamingVideoDecoder) + if not isinstance(pipeline.decoder, StreamingVideoDecoder): + raise TypeError( + "T2V requires a StreamingVideoDecoder, got " + f"{type(pipeline.decoder).__name__}." + ) ratio = pipeline.decoder.spatial_compression_ratio if scenario.pixel_height % ratio or scenario.pixel_width % ratio: raise ValueError( @@ -310,21 +396,19 @@ def __init__( ) def session_info(self) -> SessionInfo: - return SessionInfo( - output_layout="tchw", - metadata={ - FIELD_PROMPT: self.scenario.prompt, - FIELD_TOTAL_BLOCKS: self.scenario.total_blocks, - FIELD_PIXEL_HEIGHT: self.scenario.pixel_height, - FIELD_PIXEL_WIDTH: self.scenario.pixel_width, - FIELD_FPS: self.scenario.fps, - }, - ) + return SessionInfo(output_layout="tchw", metadata=self.scenario.to_mapping()) - def next_step_request(self) -> StepRequest | None: + def next_step_requirements(self) -> StepRequirements | None: if self._closed or self._step_index >= self.scenario.total_blocks: return None - return StepRequest(step_index=self._step_index) + return StepRequirements(step_index=self._step_index) + + def next_step_request(self) -> StepRequest | None: + """Legacy seam still consumed directly by the WebRTC session manager.""" + requirements = self.next_step_requirements() + if requirements is None: + return None + return StepRequest(step_index=requirements.step_index) def step(self, inputs: InferenceInput) -> StepResult: del inputs @@ -340,14 +424,12 @@ def step(self, inputs: InferenceInput) -> StepResult: metrics=stats, metadata={FIELD_PROMPT: self.scenario.prompt}, ) - self._artifact_output.write(result) + if self._artifact_output is not None: + self._artifact_output.write(result) return result def reset(self, inputs: InferenceInput | None = None) -> None: - if inputs is not None and _scenario_from_inputs(inputs) != self.scenario: - raise ValueError( - "Create a new T2V session to change the prompt or dimensions." - ) + del inputs raise RuntimeError( "T2V sessions are finite; create a new session instead of reset()." ) @@ -356,8 +438,9 @@ def close(self) -> None: if self._closed: return self._closed = True - artifacts = self._artifact_output.close() - if artifacts: + if self._artifact_output is None or self._artifact_path is None: + return + if self._artifact_output.close(): self._runtime.record_artifact(self._artifact_path, self.scenario) @@ -376,7 +459,12 @@ def create_t2v_application( input_mode=input_mode, output=output, ) - return DemoAdapterApplication(adapter=T2VDemoAdapter(model=model), spec=spec) + return DemoAdapterApplication( + adapter=T2VDemoAdapter( + model=model, write_download_artifact=output.mode == "webrtc" + ), + spec=spec, + ) def model_config_from_runner( @@ -456,48 +544,9 @@ def t2v_scenario_mapping( ) -> dict[str, object]: """Return runtime scenario values after applying launch overrides.""" defaults = defaults or T2VRunDefaults() - return { - FIELD_PROMPT: model.prompt if defaults.prompt is None else defaults.prompt, - FIELD_TOTAL_BLOCKS: ( - model.total_blocks - if defaults.total_blocks is None - else defaults.total_blocks - ), - FIELD_PIXEL_HEIGHT: ( - model.pixel_height - if defaults.pixel_height is None - else defaults.pixel_height - ), - FIELD_PIXEL_WIDTH: ( - model.pixel_width if defaults.pixel_width is None else defaults.pixel_width - ), - FIELD_FPS: model.fps if defaults.fps is None else defaults.fps, - } - - -def _scenario_from_value(value: Any, model: T2VModelConfig) -> T2VScenario: - source = value if isinstance(value, dict) else {} - prompt = str(source.get(FIELD_PROMPT, model.prompt)).strip() - if not prompt: - raise ValueError("A non-empty text-to-video prompt is required.") - return T2VScenario( - prompt=prompt, - total_blocks=_int_value(source.get(FIELD_TOTAL_BLOCKS, model.total_blocks)), - pixel_height=_int_value(source.get(FIELD_PIXEL_HEIGHT, model.pixel_height)), - pixel_width=_int_value(source.get(FIELD_PIXEL_WIDTH, model.pixel_width)), - fps=_int_value(source.get(FIELD_FPS, model.fps)), - ) - - -def _scenario_from_inputs(inputs: InferenceInput) -> T2VScenario: - source = inputs.global_conditioning - return T2VScenario( - prompt=str(source[FIELD_PROMPT]), - total_blocks=_int_value(source[FIELD_TOTAL_BLOCKS]), - pixel_height=_int_value(source[FIELD_PIXEL_HEIGHT]), - pixel_width=_int_value(source[FIELD_PIXEL_WIDTH]), - fps=_int_value(source[FIELD_FPS]), - ) + return T2VScenario.from_mapping( + defaults.scenario_overrides(), defaults=model.default_scenario() + ).to_mapping() def _int_value(value: object) -> int: diff --git a/apps/t2v/tests/test_t2v_shell.py b/apps/t2v/tests/test_t2v_shell.py index c3e4879e1..b03e6577a 100644 --- a/apps/t2v/tests/test_t2v_shell.py +++ b/apps/t2v/tests/test_t2v_shell.py @@ -4,6 +4,7 @@ from __future__ import annotations import inspect +from dataclasses import replace from types import SimpleNamespace import pytest @@ -18,7 +19,11 @@ ) from flashdreams.demo import DemoAdapterApplication -from flashdreams.runtime.demo import Mp4OutputSpec, NullOutputSpec +from flashdreams.runtime.demo import ( + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) pytestmark = pytest.mark.ci_cpu @@ -97,6 +102,35 @@ def test_t2v_shell_builds_model_config_from_runner_config() -> None: assert model.runtime_options["owner"] == "integration" +def test_t2v_shell_enables_download_artifact_only_for_webrtc() -> None: + replay_app = create_t2v_application( + model=_fake_model(), + output=Mp4OutputSpec(path="outputs/fake.mp4", fps=8, output_layout="tchw"), + ) + webrtc_app = create_t2v_application( + model=_fake_model(), input_mode="webrtc", output=WebRTCOutputSpec() + ) + + assert isinstance(replay_app, DemoAdapterApplication) + assert isinstance(webrtc_app, DemoAdapterApplication) + replay_adapter = replay_app.adapter + webrtc_adapter = webrtc_app.adapter + assert isinstance(replay_adapter, T2VDemoAdapter) + assert isinstance(webrtc_adapter, T2VDemoAdapter) + assert replay_adapter.write_download_artifact is False + assert webrtc_adapter.write_download_artifact is True + + +def test_t2v_shell_rejects_non_positive_scenario_override() -> None: + model = _fake_model() + adapter = T2VDemoAdapter(model=model) + spec = create_t2v_spec(model=model, input_mode="replay", output=NullOutputSpec()) + spec = replace(spec, scenario={**dict(spec.scenario or {}), "total_blocks": 0}) + + with pytest.raises(ValueError, match="total_blocks"): + adapter.prepare_scenario(spec) + + def test_t2v_shell_has_no_legacy_backend_imports() -> None: import t2v.t2v as t2v_shell diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index 438d16d19..6e8326a43 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -127,7 +127,11 @@ def launch_t2v( ) context = initialize_cuda_distributed(default_device=config.device) - adapter = make_adapter(defaults.backend, defaults.preset_id) + # The browser download and playback endpoints below read the session + # artifact this enables. + adapter = make_adapter( + defaults.backend, defaults.preset_id, write_download_artifact=True + ) output = WebRTCOutputSpec( host=str(host or output_overrides.get("host", "0.0.0.0")), port=_int_value( diff --git a/apps/t2v_demo/runtime.py b/apps/t2v_demo/runtime.py index 8f0040d15..410ca7fb6 100644 --- a/apps/t2v_demo/runtime.py +++ b/apps/t2v_demo/runtime.py @@ -106,9 +106,17 @@ def model_from_backend( ) -def make_adapter(backend: str, preset_id: str | None = None) -> T2VDemoAdapter: +def make_adapter( + backend: str, + preset_id: str | None = None, + *, + write_download_artifact: bool = False, +) -> T2VDemoAdapter: """Build an adapter from a legacy CLI/UI backend key.""" - return T2VDemoAdapter(model=model_from_backend(backend, preset_id)) + return T2VDemoAdapter( + model=model_from_backend(backend, preset_id), + write_download_artifact=write_download_artifact, + ) def _resolve_backend(value: str) -> T2VBackendBridge: diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 90091f635..8b47a660e 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -92,7 +92,13 @@ async def run_async(self) -> RunResult: context = self._create_context(host) result: RunResult | None = None primary_error: BaseException | None = None + app_initialized = False try: + if not host.is_healthy: + result = RunResult.rejected(reason="busy") + context.run_metrics.record_session(result) + return result + app_initialized = True self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() if isinstance(self.app, DemoAdapterApplication): @@ -116,10 +122,17 @@ async def run_async(self) -> RunResult: context=context, host=host, app=self.app, + app_initialized=app_initialized, owns_host=owns_host, run_result=result, primary_error=primary_error, - ) + ), + preserve_primary=_has_primary_outcome( + run_result=result, + primary_error=primary_error, + ), + preserved_error=primary_error + or (None if result is None else result.error), ) def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: @@ -129,7 +142,13 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: context = self._create_context(host) result: RunResult | None = None primary_error: BaseException | None = None + app_initialized = False try: + if not host.is_healthy: + result = RunResult.rejected(reason="busy") + context.run_metrics.record_session(result) + return result + app_initialized = True self.app.init(tuple(self.launch_args)) scenario = self._create_scenario() if isinstance(self.app, DemoAdapterApplication): @@ -152,6 +171,7 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: context=context, host=host, app=self.app, + app_initialized=app_initialized, owns_host=owns_host, run_result=result, primary_error=primary_error, @@ -367,19 +387,21 @@ def _close_runner_resources( context: RunContext, host: RuntimeHost, app: Application, + app_initialized: bool, owns_host: bool, run_result: RunResult | None, primary_error: BaseException | None, ) -> None: errors: list[Exception] = [] _record_cleanup_error(errors, context.close) - _close_application(errors=errors, host=host, app=app) + if app_initialized: + _close_application(errors=errors, host=host, app=app) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: _record_cleanup_notes(primary_error, errors) return - if run_result is not None and run_result.status == "failed": + if run_result is not None and run_result.status != "completed": _record_cleanup_notes(run_result.error, errors) return _raise_first_cleanup_error(errors) @@ -390,6 +412,7 @@ async def _close_runner_resources_async( context: RunContext, host: RuntimeHost, app: Application, + app_initialized: bool, owns_host: bool, run_result: RunResult | None, primary_error: BaseException | None, @@ -399,19 +422,25 @@ async def _close_runner_resources_async( await context.close_async() except Exception as exc: errors.append(exc) - await _close_application_async(errors=errors, host=host, app=app) + if app_initialized: + await _close_application_async(errors=errors, host=host, app=app) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: _record_cleanup_notes(primary_error, errors) return - if run_result is not None and run_result.status == "failed": + if run_result is not None and run_result.status != "completed": _record_cleanup_notes(run_result.error, errors) return _raise_first_cleanup_error(errors) -async def _await_runner_cleanup(cleanup: Coroutine[Any, Any, None]) -> None: +async def _await_runner_cleanup( + cleanup: Coroutine[Any, Any, None], + *, + preserve_primary: bool = False, + preserved_error: BaseException | None = None, +) -> None: cleanup_task = asyncio.create_task(cleanup) was_cancelled = False cleanup_error: BaseException | None = None @@ -431,10 +460,16 @@ async def _await_runner_cleanup(cleanup: Coroutine[Any, Any, None]) -> None: break if was_cancelled: + if preserve_primary: + _record_cleanup_notes_for_preserved_outcome(preserved_error, cleanup_error) + return cancellation = asyncio.CancelledError("cancelled during runner cleanup") _record_cancelled_cleanup_note(cancellation, cleanup_error) raise cancellation from None if cleanup_error is not None: + if preserve_primary: + _record_cleanup_notes_for_preserved_outcome(preserved_error, cleanup_error) + return raise cleanup_error @@ -467,7 +502,7 @@ def close_app() -> None: host.call(close_app) except Exception as exc: if not invoked and host.is_closed: - _record_cleanup_error(errors, app.close) + errors.append(_closed_host_cleanup_error(exc)) return errors.append(exc) @@ -489,7 +524,7 @@ def close_app() -> None: await host.call_async(close_app) except Exception as exc: if not invoked and host.is_closed: - _record_cleanup_error(errors, app.close) + errors.append(_closed_host_cleanup_error(exc)) return errors.append(exc) @@ -504,7 +539,7 @@ def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: def _record_cleanup_notes( primary: BaseException | None, - errors: Sequence[Exception], + errors: Sequence[BaseException], ) -> None: if primary is None: return @@ -525,6 +560,34 @@ def _record_cancelled_cleanup_note( add_note(f"Cleanup failed: {cleanup_error!r}") +def _record_cleanup_notes_for_preserved_outcome( + preserved_error: BaseException | None, + cleanup_error: BaseException | None, +) -> None: + if cleanup_error is not None: + _record_cleanup_notes(preserved_error, (cleanup_error,)) + + +def _has_primary_outcome( + *, + run_result: RunResult | None, + primary_error: BaseException | None, +) -> bool: + if primary_error is not None: + return True + return run_result is not None and run_result.status != "completed" + + +def _closed_host_cleanup_error(exc: Exception) -> RuntimeError: + try: + raise RuntimeError( + "Application cleanup skipped because the RuntimeHost is closed; " + "app.close() must run on the model worker." + ) from exc + except RuntimeError as cleanup_error: + return cleanup_error + + def _application_adapter(app: Application) -> DemoAdapter | None: if isinstance(app, DemoAdapterApplication): return app.adapter diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 54fa2d44b..21e470dc9 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -358,7 +358,7 @@ def test_runner_closes_public_app_when_host_is_external() -> None: assert app.closed -def test_runner_closes_public_app_when_external_host_is_closed() -> None: +def test_runner_skips_direct_app_close_for_closed_host() -> None: app = _RunnerFakeApplication(total_steps=1) host = RuntimeHost(_ExternalApplicationRuntime(app)) host.close() @@ -373,7 +373,8 @@ def test_runner_closes_public_app_when_external_host_is_closed() -> None: assert result.status == "rejected" assert result.reason == "busy" assert result.error is None - assert app.closed + assert app.init_thread_id is None + assert not app.closed def test_runner_closes_public_app_when_context_cleanup_fails() -> None: @@ -485,9 +486,7 @@ async def test_runner_run_async_delegates_to_async_session_helper() -> None: @pytest.mark.asyncio -async def test_runner_run_async_closes_public_app_when_external_host_is_closed() -> ( - None -): +async def test_runner_run_async_skips_direct_app_close_for_closed_host() -> None: app = _RunnerFakeApplication(total_steps=1) host = RuntimeHost(_ExternalApplicationRuntime(app)) host.close() @@ -503,7 +502,8 @@ async def test_runner_run_async_closes_public_app_when_external_host_is_closed() assert result.status == "rejected" assert result.reason == "busy" assert result.error is None - assert app.closed + assert app.init_thread_id is None + assert not app.closed @pytest.mark.asyncio @@ -649,6 +649,39 @@ def __init__(self, runtime: InferenceRuntime) -> None: assert created_hosts[0].is_closed +@pytest.mark.asyncio +async def test_runner_run_async_preserves_primary_error_after_cleanup_cancel() -> None: + close_started = threading.Event() + release_close = threading.Event() + app = _BlockingCloseRunnerFakeApplication( + total_steps=1, + close_started=close_started, + release_close=release_close, + ) + app.fail_init = True + runner_task = asyncio.create_task( + Runner( + io_handler=_RecordingIOHandler(), + app=app, + run_mode=_AsyncRecordingRunMode(_RecordingIOHandler()), + model_id="fake-runner", + ).run_async() + ) + + try: + assert await asyncio.to_thread(close_started.wait, 2.0) + runner_task.cancel() + await asyncio.sleep(0) + assert not runner_task.done() + finally: + release_close.set() + + with pytest.raises(RuntimeError, match="fake init failed"): + await asyncio.wait_for(runner_task, timeout=2.0) + + assert app.closed + + def test_replay_io_factory_runs_through_public_runner() -> None: output_sink = NullOutputSink(store_results=True) io_handler = create_replay_io_handler(output_sink=output_sink) From 571dd0212df13cef521f1a3dae86b6ba7b22536f Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 11:27:47 +0000 Subject: [PATCH 23/33] Move T2V replay launch into shared demo helper --- apps/t2v/__init__.py | 2 -- apps/t2v/t2v.py | 30 ------------------- apps/t2v_demo/app.py | 17 +++++++---- apps/t2v_demo/tests/test_runner.py | 28 +++++++---------- flashdreams/flashdreams/demo/__init__.py | 2 ++ flashdreams/flashdreams/demo/app.py | 19 ++++++++++-- .../tests/test_demo_application_api.py | 22 ++++++++++++++ 7 files changed, 61 insertions(+), 59 deletions(-) diff --git a/apps/t2v/__init__.py b/apps/t2v/__init__.py index 6bde43ec6..12ecf05d3 100644 --- a/apps/t2v/__init__.py +++ b/apps/t2v/__init__.py @@ -19,7 +19,6 @@ create_t2v_application, create_t2v_spec, model_config_from_runner, - run_t2v_replay_application, t2v_scenario_mapping, ) @@ -39,6 +38,5 @@ "create_t2v_application", "create_t2v_spec", "model_config_from_runner", - "run_t2v_replay_application", "t2v_scenario_mapping", ] diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index 464d42480..bc2048720 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -16,9 +16,6 @@ from flashdreams.demo import ( Application, DemoAdapterApplication, - FileOutputSink, - Runner, - create_replay_io_handler, ) from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.video_output import VideoOutputStream @@ -36,13 +33,11 @@ from flashdreams.runtime._utils import freeze_mapping from flashdreams.runtime.demo import ( DemoSpec, - Mp4OutputSpec, NullOutputSpec, OutputSpec, PreparedScenario, ) from flashdreams.runtime.demo.outputs import SessionInfo -from flashdreams.runtime.demo.run_modes import RunResult from flashdreams.runtime.demo.session_inputs import ( PreparedStep, ProviderCapabilities, @@ -515,30 +510,6 @@ def create_t2v_spec( ) -def run_t2v_replay_application( - *, - model: T2VModelConfig, - defaults: T2VRunDefaults | None, - output: Mp4OutputSpec | NullOutputSpec, -) -> RunResult: - """Run finite T2V replay through the public runner and replay IO handler.""" - output_sink = None - if isinstance(output, Mp4OutputSpec): - output_sink = FileOutputSink( - output_path=Path(output.path), - fps=output.fps, - output_layout=output.output_layout, - ) - result = Runner( - io_handler=create_replay_io_handler(output_sink=output_sink), - app=create_t2v_application(model=model, defaults=defaults, output=output), - ).run() - if result.status != "completed": - reason = result.reason or str(result.error) or "T2V replay failed." - raise RuntimeError(reason) - return result - - def t2v_scenario_mapping( *, model: T2VModelConfig, defaults: T2VRunDefaults | None = None ) -> dict[str, object]: @@ -592,6 +563,5 @@ def _validate_optional_positive_int(value: int | None, *, name: str) -> None: "create_t2v_application", "create_t2v_spec", "model_config_from_runner", - "run_t2v_replay_application", "t2v_scenario_mapping", ] diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py index 6e8326a43..498374d99 100644 --- a/apps/t2v_demo/app.py +++ b/apps/t2v_demo/app.py @@ -18,11 +18,10 @@ T2VRunDefaults, create_t2v_application, create_t2v_spec, - run_t2v_replay_application, t2v_scenario_mapping, ) -from flashdreams.demo import Application +from flashdreams.demo import Application, run_application_replay from flashdreams.runtime.demo import ( DemoSpec, Mp4OutputSpec, @@ -120,11 +119,17 @@ def launch_t2v( output_overrides.get("fps", scenario[FIELD_FPS]), name="fps" ), ) - return run_t2v_replay_application( - model=model_from_backend(defaults.backend, defaults.preset_id), - defaults=defaults, - output=output, + result = run_application_replay( + app=create_t2v_application( + model=model_from_backend(defaults.backend, defaults.preset_id), + defaults=defaults, + output=output, + ) ) + if result.status != "completed": + reason = result.reason or str(result.error) or "T2V replay failed." + raise RuntimeError(reason) + return result context = initialize_cuda_distributed(default_device=config.device) # The browser download and playback endpoints below read the session diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py index 1cf6da22c..295ff50de 100644 --- a/apps/t2v_demo/tests/test_runner.py +++ b/apps/t2v_demo/tests/test_runner.py @@ -15,10 +15,8 @@ from flashdreams.demo import ( Application, DemoAdapterApplication, - FileOutputSink, - ReplayIOHandler, ) -from flashdreams.runtime.demo import RunResult +from flashdreams.runtime.demo import Mp4OutputSpec, RunResult pytestmark = pytest.mark.ci_cpu @@ -103,15 +101,11 @@ def test_runner_mp4_launch_uses_demo_entrypoint( ) -> None: captured: dict[str, object] = {} - class FakeRunner: - def __init__(self, *, io_handler: object, app: object) -> None: - captured["io_handler"] = io_handler - captured["app"] = app + def fake_run_application_replay(*, app: Application) -> RunResult: + captured["app"] = app + return RunResult(status="completed") - def run(self) -> RunResult: - return RunResult(status="completed") - - monkeypatch.setattr(t2v_shell, "Runner", FakeRunner) + monkeypatch.setattr(app, "run_application_replay", fake_run_application_replay) config = T2VDemoRunnerConfig( runner_name="t2v", description="test", @@ -127,15 +121,13 @@ def run(self) -> RunResult: ) public_app = captured["app"] - io_handler = captured["io_handler"] assert isinstance(public_app, DemoAdapterApplication) - assert isinstance(io_handler, ReplayIOHandler) scenario = public_app.spec.scenario assert isinstance(scenario, dict) assert scenario["prompt"] == "A waterfall" assert scenario["total_blocks"] == 3 - output_sink = io_handler.output_sink - assert isinstance(output_sink, FileOutputSink) - assert str(output_sink.output_path) == "outputs/test.mp4" - assert output_sink.fps == 24 - assert output_sink.output_layout == "tchw" + output = public_app.spec.output + assert isinstance(output, Mp4OutputSpec) + assert str(output.path) == "outputs/test.mp4" + assert output.fps == 24 + assert output.output_layout == "tchw" diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 6e2005b0a..776aa7b32 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -6,6 +6,7 @@ from flashdreams.demo.app import ( DemoApplication, create_demo_application, + run_application_replay, run_replay_application, ) from flashdreams.demo.application import ( @@ -82,5 +83,6 @@ "create_replay_io_handler", "create_webrtc_io_handler", "input_state_from_window", + "run_application_replay", "run_replay_application", ] diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index 9b1b9bd2e..3dd444e61 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -127,9 +127,17 @@ def _run_handler(self, args: argparse.Namespace, handler: IOHandler) -> RunResul def run_replay_application(*, spec: DemoSpec, adapter: DemoAdapter) -> RunResult: """Run a finite demo spec through the public replay IO factory and runner.""" + return run_application_replay(app=DemoAdapterApplication(adapter=adapter, spec=spec)) + + +def run_application_replay(*, app: Application) -> RunResult: + """Run a finite public application through the replay IO factory and runner.""" + output_sink = None + if isinstance(app, DemoAdapterApplication): + output_sink = build_output_sink(app.spec.output) return Runner( - io_handler=create_replay_io_handler(output_sink=build_output_sink(spec.output)), - app=DemoAdapterApplication(adapter=adapter, spec=spec), + io_handler=create_replay_io_handler(output_sink=output_sink), + app=app, ).run() @@ -159,4 +167,9 @@ def _raise_for_failed_result(result: RunResult) -> None: raise SystemExit(1) -__all__ = ["DemoApplication", "create_demo_application", "run_replay_application"] +__all__ = [ + "DemoApplication", + "create_demo_application", + "run_application_replay", + "run_replay_application", +] diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 21e470dc9..3821bbb6b 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -38,6 +38,7 @@ create_replay_io_handler, create_webrtc_io_handler, input_state_from_window, + run_application_replay, ) from flashdreams.runtime import ( CanonicalInputSchema, @@ -93,6 +94,7 @@ def test_public_demo_contracts_are_importable() -> None: assert IOHandler.__name__ == "IOHandler" assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" + assert run_application_replay.__name__ == "run_application_replay" assert create_demo_application.__name__ == "create_demo_application" assert create_replay_io_handler.__name__ == "create_replay_io_handler" assert create_native_window_io_handler.__name__ == ( @@ -809,6 +811,26 @@ def test_demo_application_replay_selects_factory_and_runner() -> None: assert runtime.closed +def test_run_application_replay_uses_demo_adapter_spec() -> None: + adapter = _FakeDemoAdapter() + app = DemoAdapterApplication( + adapter=adapter, + spec=DemoSpec( + model_id="fake-demo", + input_mode="replay", + output=NullOutputSpec(), + ), + ) + + result = run_application_replay(app=app) + + assert result.status == "completed" + assert adapter.runtimes + runtime = adapter.runtimes[0] + assert runtime.session.closed + assert runtime.closed + + def test_demo_application_can_be_built_from_callbacks() -> None: adapter = _FakeDemoAdapter() app = create_demo_application( From c2916faff924bf566299243b8a3d7b797ecebd04 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 11:36:51 +0000 Subject: [PATCH 24/33] Close public apps from RuntimeHost shutdown --- flashdreams/flashdreams/demo/runner.py | 62 ++++++++++++---- flashdreams/flashdreams/runtime/demo/host.py | 59 ++++++++++++++- .../tests/test_demo_application_api.py | 73 +++++++++++++++++++ 3 files changed, 179 insertions(+), 15 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 8b47a660e..d49c36fd2 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -7,7 +7,7 @@ import asyncio import inspect -from collections.abc import Coroutine, Sequence +from collections.abc import Callable, Coroutine, Sequence from dataclasses import dataclass, field from typing import Any, cast @@ -93,13 +93,17 @@ async def run_async(self) -> RunResult: result: RunResult | None = None primary_error: BaseException | None = None app_initialized = False + app_cleanup: _ApplicationCleanup | None = None + remove_app_close_hook: Callable[[], None] | None = None try: if not host.is_healthy: result = RunResult.rejected(reason="busy") context.run_metrics.record_session(result) return result + app_cleanup = _ApplicationCleanup(self.app) app_initialized = True self.app.init(tuple(self.launch_args)) + remove_app_close_hook = host.add_close_hook(app_cleanup.close) scenario = self._create_scenario() if isinstance(self.app, DemoAdapterApplication): _configure_replay_io_handler(self.io_handler, scenario) @@ -121,8 +125,9 @@ async def run_async(self) -> RunResult: _close_runner_resources_async( context=context, host=host, - app=self.app, + app_cleanup=app_cleanup, app_initialized=app_initialized, + remove_app_close_hook=remove_app_close_hook, owns_host=owns_host, run_result=result, primary_error=primary_error, @@ -143,13 +148,17 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: result: RunResult | None = None primary_error: BaseException | None = None app_initialized = False + app_cleanup: _ApplicationCleanup | None = None + remove_app_close_hook: Callable[[], None] | None = None try: if not host.is_healthy: result = RunResult.rejected(reason="busy") context.run_metrics.record_session(result) return result + app_cleanup = _ApplicationCleanup(self.app) app_initialized = True self.app.init(tuple(self.launch_args)) + remove_app_close_hook = host.add_close_hook(app_cleanup.close) scenario = self._create_scenario() if isinstance(self.app, DemoAdapterApplication): _configure_replay_io_handler(self.io_handler, scenario) @@ -170,8 +179,9 @@ def _run_sync(self, *, run_mode: RunMode | None = None) -> RunResult: _close_runner_resources( context=context, host=host, - app=self.app, + app_cleanup=app_cleanup, app_initialized=app_initialized, + remove_app_close_hook=remove_app_close_hook, owns_host=owns_host, run_result=result, primary_error=primary_error, @@ -355,6 +365,18 @@ def close(self) -> None: self.closed = True +@dataclass(slots=True) +class _ApplicationCleanup: + app: Application + closed: bool = False + + def close(self) -> None: + if self.closed: + return + self.app.close() + self.closed = True + + def _run_mode_is_async(run_mode: RunMode) -> bool: driver = run_mode.select_driver() return inspect.iscoroutinefunction(driver.run_one_session) @@ -386,16 +408,19 @@ def _close_runner_resources( *, context: RunContext, host: RuntimeHost, - app: Application, + app_cleanup: _ApplicationCleanup | None, app_initialized: bool, + remove_app_close_hook: Callable[[], None] | None, owns_host: bool, run_result: RunResult | None, primary_error: BaseException | None, ) -> None: errors: list[Exception] = [] _record_cleanup_error(errors, context.close) - if app_initialized: - _close_application(errors=errors, host=host, app=app) + if app_initialized and app_cleanup is not None: + _close_application(errors=errors, host=host, cleanup=app_cleanup) + if remove_app_close_hook is not None: + _record_cleanup_error(errors, remove_app_close_hook) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: @@ -411,8 +436,9 @@ async def _close_runner_resources_async( *, context: RunContext, host: RuntimeHost, - app: Application, + app_cleanup: _ApplicationCleanup | None, app_initialized: bool, + remove_app_close_hook: Callable[[], None] | None, owns_host: bool, run_result: RunResult | None, primary_error: BaseException | None, @@ -422,8 +448,10 @@ async def _close_runner_resources_async( await context.close_async() except Exception as exc: errors.append(exc) - if app_initialized: - await _close_application_async(errors=errors, host=host, app=app) + if app_initialized and app_cleanup is not None: + await _close_application_async(errors=errors, host=host, cleanup=app_cleanup) + if remove_app_close_hook is not None: + _record_cleanup_error(errors, remove_app_close_hook) if owns_host: _record_cleanup_error(errors, host.close) if primary_error is not None: @@ -489,18 +517,22 @@ def _close_application( *, errors: list[Exception], host: RuntimeHost, - app: Application, + cleanup: _ApplicationCleanup, ) -> None: + if cleanup.closed: + return invoked = False def close_app() -> None: nonlocal invoked invoked = True - app.close() + cleanup.close() try: host.call(close_app) except Exception as exc: + if cleanup.closed: + return if not invoked and host.is_closed: errors.append(_closed_host_cleanup_error(exc)) return @@ -511,18 +543,22 @@ async def _close_application_async( *, errors: list[Exception], host: RuntimeHost, - app: Application, + cleanup: _ApplicationCleanup, ) -> None: + if cleanup.closed: + return invoked = False def close_app() -> None: nonlocal invoked invoked = True - app.close() + cleanup.close() try: await host.call_async(close_app) except Exception as exc: + if cleanup.closed: + return if not invoked and host.is_closed: errors.append(_closed_host_cleanup_error(exc)) return diff --git a/flashdreams/flashdreams/runtime/demo/host.py b/flashdreams/flashdreams/runtime/demo/host.py index 787e07147..90076e570 100644 --- a/flashdreams/flashdreams/runtime/demo/host.py +++ b/flashdreams/flashdreams/runtime/demo/host.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field +from threading import Lock from typing import TypeVar from flashdreams.runtime._utils import freeze_mapping @@ -58,6 +59,8 @@ def __init__( self._worker_loop = worker_loop self._healthy = True self._closed = False + self._close_hooks: list[Callable[[], None]] = [] + self._close_hooks_lock = Lock() self._unhealthy_reason: str | None = None self._unhealthy_error: Exception | None = None @@ -140,6 +143,21 @@ async def call_async( self._require_open() return await self._worker.call(func, *args, **kwargs) + def add_close_hook(self, hook: Callable[[], None]) -> Callable[[], None]: + """Run ``hook`` on the model worker before the hosted runtime closes.""" + with self._close_hooks_lock: + self._require_open() + self._close_hooks.append(hook) + + def remove_hook() -> None: + with self._close_hooks_lock: + try: + self._close_hooks.remove(hook) + except ValueError: + pass + + return remove_hook + def start_session(self, inputs: InferenceInput) -> InferenceSession: """Start one inference session through the hosted runtime.""" self._require_open() @@ -159,9 +177,17 @@ def close(self) -> None: """Close runtime-owned state and stop the model-execution worker.""" if self._closed: return + errors: list[Exception] = [] try: - self._worker.call_blocking(self._runtime.close) - self._call_optional_runtime_hook("close_distributed") + for hook in self._take_close_hooks(): + _record_close_error(errors, self._worker.call_blocking, hook) + _record_close_error(errors, self._worker.call_blocking, self._runtime.close) + close_distributed = getattr(self._runtime, "close_distributed", None) + if callable(close_distributed): + _record_close_error( + errors, self._worker.call_blocking, close_distributed + ) + _raise_first_close_error(errors) finally: self._closed = True self._worker.close_blocking() @@ -175,5 +201,34 @@ def _require_open(self) -> None: if self._closed: raise RuntimeError("runtime host is closed") + def _take_close_hooks(self) -> tuple[Callable[[], None], ...]: + with self._close_hooks_lock: + hooks = tuple(self._close_hooks) + self._close_hooks.clear() + return hooks + + +def _record_close_error( + errors: list[Exception], + close: Callable[..., object], + /, + *args: object, +) -> None: + try: + close(*args) + except Exception as exc: + errors.append(exc) + + +def _raise_first_close_error(errors: Sequence[Exception]) -> None: + if not errors: + return + first = errors[0] + add_note = getattr(first, "add_note", None) + for extra in errors[1:]: + if callable(add_note): + add_note(f"Additional host close error: {extra!r}") + raise first + __all__ = ["ModelWarmupPlan", "RuntimeHost", "WarmupSessionInputs"] diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 3821bbb6b..5ce4f0523 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -360,6 +360,27 @@ def test_runner_closes_public_app_when_host_is_external() -> None: assert app.closed +def test_runner_external_host_close_hook_closes_public_app() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = RuntimeHost(_ExternalApplicationRuntime(app)) + + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="host-closes-after-init", + driver=_ClosingHostDriver(), + ), + model_id="fake-runner", + ).run() + + assert result.status == "completed" + assert host.is_closed + assert app.closed + + def test_runner_skips_direct_app_close_for_closed_host() -> None: app = _RunnerFakeApplication(total_steps=1) host = RuntimeHost(_ExternalApplicationRuntime(app)) @@ -508,6 +529,28 @@ async def test_runner_run_async_skips_direct_app_close_for_closed_host() -> None assert not app.closed +@pytest.mark.asyncio +async def test_runner_run_async_external_host_close_hook_closes_public_app() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = RuntimeHost(_ExternalApplicationRuntime(app)) + + result = await Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="async-host-closes-after-init", + driver=_AsyncClosingHostDriver(), + ), + model_id="fake-runner", + ).run_async() + + assert result.status == "completed" + assert host.is_closed + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_closes_public_app_when_context_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -1469,6 +1512,21 @@ def run_one_session( ) +class _ClosingHostDriver: + def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + host.call(provider.close) + host.close() + return session_edges.close_result(status="completed") + + @dataclass(slots=True) class _AsyncRecordingRunMode: io_handler: _RecordingIOHandler @@ -1568,6 +1626,21 @@ async def run_one_session( ) +class _AsyncClosingHostDriver: + async def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + await host.call_async(provider.close) + host.close() + return session_edges.close_result(status="completed") + + class _AsyncRunModeInputSource: is_finite = False is_deterministic = False From 9d7f10884fbf27d29d7077252a3126514a5ec9c3 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 11:52:43 +0000 Subject: [PATCH 25/33] Launch public apps directly from flashdreams-run Add finite-mode application launch support to the central CLI for discovered flashdreams.applications entry points. Route direct mp4 and null launches through the current public demo Application/Runner API, while keeping the existing t2v runner bridge for compatibility. --- apps/t2v/t2v.py | 1 + flashdreams/flashdreams/demo/app.py | 7 +- flashdreams/flashdreams/scripts/cli.py | 295 +++++++++++++++++++++- flashdreams/tests/test_launch_manifest.py | 117 +++++++++ 4 files changed, 417 insertions(+), 3 deletions(-) diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index bc2048720..0bf32f111 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -507,6 +507,7 @@ def create_t2v_spec( **defaults.runtime_options, }, ), + metadata={"output_layout": "tchw"}, ) diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index 3dd444e61..f943cd7c5 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -7,7 +7,7 @@ import argparse import sys -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import Any import torch @@ -130,7 +130,9 @@ def run_replay_application(*, spec: DemoSpec, adapter: DemoAdapter) -> RunResult return run_application_replay(app=DemoAdapterApplication(adapter=adapter, spec=spec)) -def run_application_replay(*, app: Application) -> RunResult: +def run_application_replay( + *, app: Application, launch_args: Sequence[str] = () +) -> RunResult: """Run a finite public application through the replay IO factory and runner.""" output_sink = None if isinstance(app, DemoAdapterApplication): @@ -138,6 +140,7 @@ def run_application_replay(*, app: Application) -> RunResult: return Runner( io_handler=create_replay_io_handler(output_sink=output_sink), app=app, + launch_args=tuple(launch_args), ).run() diff --git a/flashdreams/flashdreams/scripts/cli.py b/flashdreams/flashdreams/scripts/cli.py index c366debcd..afb246e80 100644 --- a/flashdreams/flashdreams/scripts/cli.py +++ b/flashdreams/flashdreams/scripts/cli.py @@ -43,7 +43,7 @@ import dataclasses import os import sys -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from pathlib import Path from typing import Annotated, Any, cast @@ -53,7 +53,14 @@ from flashdreams.configs.runner_configs import _annotated_base_runner_union, all_runners from flashdreams.core.distributed import shutdown as shutdown_distributed from flashdreams.core.io.disk import disk_space_error_from_exception +from flashdreams.demo import ( + Application, + DemoAdapterApplication, + run_application_replay, +) from flashdreams.infra.runner import RunnerConfig +from flashdreams.plugins import discover_applications +from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, NullOutputSpec from flashdreams.serving.launch import ( LaunchMode, LaunchOptions, @@ -199,6 +206,20 @@ def entrypoint(argv: list[str] | None = None) -> None: """ tyro.extras.set_accent_color("bright_yellow") raw_args = list(sys.argv[1:] if argv is None else argv) + runners = dict(all_runners()) + target_name = _first_positional_arg(raw_args) + if target_name is not None and target_name not in runners: + applications = discover_applications() + application_name = _selected_application_name(target_name, applications) + if application_name is not None: + _run_with_disk_error_handling( + lambda: _entrypoint_application( + raw_args, + application_name=application_name, + application=applications[application_name], + ) + ) + return ( normalized_args, runners, @@ -322,6 +343,278 @@ def entrypoint(argv: list[str] | None = None) -> None: ) +def _first_positional_arg(args: Sequence[str]) -> str | None: + index = 0 + while index < len(args): + item = args[index] + if item in {"--no-instantiate", "--prefer-sw-encoder", "--help", "-h"}: + index += 1 + continue + if item in {"--host", "--port", "--manifest"}: + index += 2 + continue + if any(item.startswith(option + "=") for option in ("--host", "--port")): + index += 1 + continue + parsed_override = _parse_launch_override_token(item) + if parsed_override is not None: + _section, _key, inline_value = parsed_override + index += 1 if inline_value is not None else 2 + continue + if item.startswith("-"): + index += 1 + continue + return item + return None + + +def _selected_application_name( + target_name: str, + applications: Mapping[str, Application], +) -> str | None: + if target_name in applications: + return target_name + return None + + +def _entrypoint_application( + raw_args: list[str], + *, + application_name: str, + application: Application, +) -> None: + if "--help" in raw_args or "-h" in raw_args: + _print_application_help(application_name, application) + raise SystemExit(0) + mode, launch_args, no_instantiate, launch_overrides = ( + _prepare_application_cli_args(raw_args, application_name=application_name) + ) + if mode not in {"mp4", "null"}: + raise ValueError( + f"Application {application_name!r} currently supports direct finite " + "launch modes 'mp4' and 'null'. Use a compatibility runner for " + f"{mode!r}." + ) + scenario = dict(launch_overrides.scenario) + output = dict(launch_overrides.output) + configured = _configure_application_replay( + application=application, + application_name=application_name, + mode=mode, + scenario_overrides=scenario, + output_overrides=output, + ) + if _is_rank_zero(): + print(f"Resolved application: {application_name!r}") + print(f"Launch mode: {mode}") + if scenario: + print(f"Scenario: {scenario}") + if output: + print(f"Output settings: {output}") + if no_instantiate: + return + _handle_launch_result( + run_application_replay(app=configured, launch_args=launch_args) + ) + + +def _prepare_application_cli_args( + args: list[str], + *, + application_name: str, +) -> tuple[LaunchMode, tuple[str, ...], bool, _LaunchCliOverrides]: + normalized, launch_overrides = _pop_launch_overrides(args) + normalized, manifest_path = _pop_option(normalized, "--manifest") + if manifest_path is not None: + raise ValueError( + "--manifest is not supported for direct application launches yet; " + "use --scenario.KEY and --output.KEY overrides." + ) + normalized = _hoist_global_options(normalized) + no_instantiate = False + remaining: list[str] = [] + index = 0 + while index < len(normalized): + item = normalized[index] + if item == "--no-instantiate": + no_instantiate = True + index += 1 + continue + if item == "--prefer-sw-encoder": + raise ValueError( + "--prefer-sw-encoder is only supported by WebRTC runner launches." + ) + if item in {"--host", "--port"}: + raise ValueError(f"{item} is only supported by WebRTC runner launches.") + if any(item.startswith(option + "=") for option in ("--host", "--port")): + raise ValueError( + f"{item.split('=', 1)[0]} is only supported by WebRTC runner launches." + ) + remaining.append(item) + index += 1 + + try: + app_index = remaining.index(application_name) + except ValueError as exc: + raise ValueError( + f"Application slug {application_name!r} was not present in argv." + ) from exc + del remaining[app_index] + raw_mode: LaunchMode = "run" + if app_index < len(remaining) and remaining[app_index] in _POSITIONAL_MODES: + raw_mode = cast(LaunchMode, remaining.pop(app_index)) + return raw_mode, tuple(remaining), no_instantiate, launch_overrides + + +def _configure_application_replay( + *, + application: Application, + application_name: str, + mode: LaunchMode, + scenario_overrides: Mapping[str, object], + output_overrides: Mapping[str, object], +) -> Application: + if not isinstance(application, DemoAdapterApplication): + if scenario_overrides or output_overrides or mode != "null": + raise ValueError( + "Direct application launch with scenario/output overrides requires " + "a DemoAdapterApplication." + ) + return application + + scenario = _application_scenario(application.spec, scenario_overrides) + output = _application_output_spec( + application_name=application_name, + mode=mode, + spec=application.spec, + scenario=scenario, + output_overrides=output_overrides, + ) + return DemoAdapterApplication( + adapter=application.adapter, + spec=dataclasses.replace( + application.spec, + input_mode="replay", + scenario=scenario, + output=output, + ), + ) + + +def _application_scenario( + spec: DemoSpec, + overrides: Mapping[str, object], +) -> object: + if not overrides: + return spec.scenario + if spec.scenario is None: + return dict(overrides) + if not isinstance(spec.scenario, Mapping): + raise ValueError( + "Scenario overrides require the application DemoSpec.scenario to be " + f"a mapping, got {type(spec.scenario).__name__}." + ) + return {**dict(spec.scenario), **dict(overrides)} + + +def _application_output_spec( + *, + application_name: str, + mode: LaunchMode, + spec: DemoSpec, + scenario: object, + output_overrides: Mapping[str, object], +) -> Mp4OutputSpec | NullOutputSpec: + if mode == "null": + _reject_unknown_output_keys(output_overrides, allowed={"store_results"}) + return NullOutputSpec( + store_results=bool(output_overrides.get("store_results", False)) + ) + if mode != "mp4": + raise ValueError( + f"Direct application launch mode {mode!r} is not implemented." + ) + _reject_unknown_output_keys( + output_overrides, + allowed={"fps", "layout", "move_to_cpu", "output", "output_layout", "path"}, + ) + current_output = spec.output + path = output_overrides.get("path", output_overrides.get("output")) + if path is None and isinstance(current_output, Mp4OutputSpec): + path = current_output.path + if path is None: + path = Path("outputs") / f"{application_name}.mp4" + fps = output_overrides.get("fps") + if fps is None and isinstance(current_output, Mp4OutputSpec): + fps = current_output.fps + if fps is None: + fps = _scenario_field(scenario, "fps") + if fps is None: + fps = spec.metadata.get("fps") + if fps is None: + raise ValueError( + "Direct application MP4 launch requires --output.fps or an fps " + "value in the application scenario or metadata." + ) + output_layout = output_overrides.get( + "output_layout", + output_overrides.get("layout"), + ) + if output_layout is None and isinstance(current_output, Mp4OutputSpec): + output_layout = current_output.output_layout + if output_layout is None: + output_layout = spec.metadata.get("output_layout", "bvtchw") + move_to_cpu = output_overrides.get("move_to_cpu") + if move_to_cpu is None and isinstance(current_output, Mp4OutputSpec): + move_to_cpu = current_output.move_to_cpu + return Mp4OutputSpec( + path=Path(str(path)), + fps=_positive_number(fps, name="fps"), + output_layout=cast(Any, str(output_layout)), + move_to_cpu=bool(True if move_to_cpu is None else move_to_cpu), + ) + + +def _scenario_field(scenario: object, name: str) -> object | None: + if isinstance(scenario, Mapping): + return scenario.get(name) + return None + + +def _positive_number(value: object, *, name: str) -> int | float: + if isinstance(value, bool): + raise TypeError(f"{name} must be numeric, not bool.") + if isinstance(value, int | float): + number = value + elif isinstance(value, str): + number = float(value) if "." in value else int(value) + else: + raise TypeError(f"{name} must be numeric, got {type(value).__name__}.") + if float(number) <= 0: + raise ValueError(f"{name} must be > 0.") + return number + + +def _reject_unknown_output_keys( + output_overrides: Mapping[str, object], + *, + allowed: set[str], +) -> None: + unknown = sorted(set(output_overrides) - allowed) + if unknown: + raise ValueError(f"Unsupported application output fields: {', '.join(unknown)}") + + +def _print_application_help(application_name: str, application: Application) -> None: + modes = ("null",) + if isinstance(application, DemoAdapterApplication): + supported = set(application.adapter.supported_output_modes()) + modes = tuple(mode for mode in ("mp4", "null") if mode in supported) + print(f"Usage: flashdreams-run {application_name} [options]") + print(f"Available direct application modes: {', '.join(modes)}") + print("Use --scenario.KEY VALUE and --output.KEY VALUE for finite-mode settings.") + + def _prepare_cli_args( args: list[str], ) -> tuple[ diff --git a/flashdreams/tests/test_launch_manifest.py b/flashdreams/tests/test_launch_manifest.py index 7adbdd152..bd572c3cb 100644 --- a/flashdreams/tests/test_launch_manifest.py +++ b/flashdreams/tests/test_launch_manifest.py @@ -9,7 +9,23 @@ import pytest +from flashdreams.demo import DemoAdapterApplication from flashdreams.infra.runner import RunnerConfig +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InputMapping, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + PreparedScenario, + RunResult, +) from flashdreams.scripts import cli from flashdreams.serving.launch import ResolvedLaunch, resolve_launch from flashdreams.serving.launch_manifest import load_launch_manifest @@ -29,6 +45,20 @@ def _config(name: str = "demo-runner") -> RunnerConfig: ) +def _application() -> DemoAdapterApplication: + return DemoAdapterApplication( + adapter=_CliApplicationAdapter(), + spec=DemoSpec( + model_id="test-app", + input_mode="replay", + scenario={"prompt": "default", "fps": 16}, + output=NullOutputSpec(), + config=InferenceConfig(model_id="test-app"), + metadata={"output_layout": "tchw"}, + ), + ) + + def test_launch_manifest_loads_strict_sections_and_relative_paths( tmp_path: Path, ) -> None: @@ -57,6 +87,67 @@ def test_launch_manifest_loads_strict_sections_and_relative_paths( assert manifest.apply_runner_overrides(_config()).device == "cuda:3" +def test_entrypoint_launches_application_slug_mp4_directly( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: list[tuple[DemoAdapterApplication, tuple[str, ...]]] = [] + + def fake_run_application_replay( + *, app: object, launch_args: tuple[str, ...] = () + ) -> RunResult: + assert isinstance(app, DemoAdapterApplication) + captured.append((app, launch_args)) + return RunResult(status="completed") + + monkeypatch.setattr(cli, "all_runners", lambda: {}) + monkeypatch.setattr( + cli, + "discover_applications", + lambda: {"test-app": _application()}, + ) + monkeypatch.setattr(cli, "run_application_replay", fake_run_application_replay) + + cli.entrypoint( + [ + "test-app", + "mp4", + "--scenario.prompt", + "A waterfall", + "--scenario.fps", + "12", + "--output.path", + str(tmp_path / "test.mp4"), + ] + ) + + configured, launch_args = captured[0] + assert launch_args == () + assert configured.spec.scenario == { + "prompt": "A waterfall", + "fps": 12, + } + output = configured.spec.output + assert isinstance(output, Mp4OutputSpec) + assert output.path == tmp_path / "test.mp4" + assert output.fps == 12 + assert output.output_layout == "tchw" + + +def test_entrypoint_application_null_rejects_output_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(cli, "all_runners", lambda: {}) + monkeypatch.setattr( + cli, + "discover_applications", + lambda: {"test-app": _application()}, + ) + + with pytest.raises(ValueError, match="Unsupported application output fields: path"): + cli.entrypoint(["test-app", "null", "--output.path", "unexpected.mp4"]) + + def test_launch_manifest_does_not_guess_configs_directory( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -441,3 +532,29 @@ def test_documented_launch_manifests_resolve(filename: str) -> None: ), ) assert resolved.mode == manifest.mode + + +class _CliApplicationAdapter: + model_id = "test-app" + inference_input_schema = InferenceInputSchema() + canonical_input_schema = CanonicalInputSchema() + + def supported_input_modes(self) -> tuple[str, ...]: + return ("replay",) + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "null") + + def default_input_mapping(self) -> InputMapping: + return IdentityInputMapping() + + def validate_config(self, config: InferenceConfig) -> None: + assert config.model_id == self.model_id + + def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: + del spec + return PreparedScenario(initial_inputs=InferenceInput()) + + def create_runtime(self, config: InferenceConfig) -> object: + del config + raise AssertionError("direct app CLI test should not instantiate the runtime") From b9e80abc5d396a26b542ab9ca18f06403e911c80 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 12:12:53 +0000 Subject: [PATCH 26/33] Add direct WebRTC launch for public T2V apps Route discovered application slugs through the shared WebRTC serving path for direct flashdreams-run app launches. Move T2V WebRTC resources into the neutral T2V shell and keep the legacy t2v_demo package until the new WebRTC path is remotely validated. --- apps/t2v/pyproject.toml | 5 +- apps/t2v/t2v.py | 142 +++++++++- apps/t2v/web/adapter.css | 3 + apps/t2v/web/adapter.js | 23 ++ flashdreams/flashdreams/demo/__init__.py | 2 + flashdreams/flashdreams/demo/app.py | 154 ++++++++++- .../flashdreams/runtime/video_runner.py | 249 +++++++++++++++++ flashdreams/flashdreams/scripts/cli.py | 260 ++++++++++++++++-- .../flashdreams/serving/webrtc/manager.py | 13 + .../tests/test_demo_application_api.py | 2 + flashdreams/tests/test_launch_manifest.py | 51 +++- .../causal_forcing/causal_forcing/config.py | 45 ++- .../causal_forcing/causal_forcing/runner.py | 256 ----------------- .../cosmos_predict2/cosmos_predict2/config.py | 47 +++- .../cosmos_predict2/cosmos_predict2/runner.py | 225 --------------- .../fastvideo_causal_wan22/config.py | 31 ++- .../fastvideo_causal_wan22/runner.py | 174 ------------ .../self_forcing/self_forcing/config.py | 29 +- .../self_forcing/self_forcing/runner.py | 172 ------------ integrations/wan21/wan21/config.py | 44 ++- integrations/wan21/wan21/runner.py | 250 ----------------- uv.lock | 4 +- 22 files changed, 1061 insertions(+), 1120 deletions(-) create mode 100644 apps/t2v/web/adapter.css create mode 100644 apps/t2v/web/adapter.js create mode 100644 flashdreams/flashdreams/runtime/video_runner.py delete mode 100644 integrations/causal_forcing/causal_forcing/runner.py delete mode 100644 integrations/cosmos_predict2/cosmos_predict2/runner.py delete mode 100644 integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py delete mode 100644 integrations/self_forcing/self_forcing/runner.py delete mode 100644 integrations/wan21/wan21/runner.py diff --git a/apps/t2v/pyproject.toml b/apps/t2v/pyproject.toml index e48918087..4521941aa 100644 --- a/apps/t2v/pyproject.toml +++ b/apps/t2v/pyproject.toml @@ -10,7 +10,7 @@ name = "flashdreams-t2v" version = "0.1.0" description = "Shared model-neutral text-to-video demo application shell" requires-python = ">=3.10" -dependencies = ["flashdreams"] +dependencies = ["flashdreams[serving]"] [tool.uv.sources] flashdreams = { workspace = true } @@ -18,3 +18,6 @@ flashdreams = { workspace = true } [tool.setuptools] packages = ["t2v"] package-dir = { t2v = "." } + +[tool.setuptools.package-data] +t2v = ["web/*.js", "web/*.css"] diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py index 0bf32f111..4b0ed5402 100644 --- a/apps/t2v/t2v.py +++ b/apps/t2v/t2v.py @@ -5,8 +5,12 @@ from __future__ import annotations +import io +import json +import zipfile from collections.abc import Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from importlib.resources import files from pathlib import Path from typing import Any from uuid import uuid4 @@ -36,6 +40,8 @@ NullOutputSpec, OutputSpec, PreparedScenario, + WebRTCAppResources, + WebRTCOutputSpec, ) from flashdreams.runtime.demo.outputs import SessionInfo from flashdreams.runtime.demo.session_inputs import ( @@ -250,6 +256,32 @@ def create_model_input_provider( del spec return T2VInputProvider(initial_inputs=scenario.initial_inputs) + def configure_for_output(self, output: OutputSpec) -> "T2VDemoAdapter": + """Return an adapter configured for output-specific T2V behavior.""" + return T2VDemoAdapter( + model=self.model, + write_download_artifact=output.mode == "webrtc", + ) + + def create_webrtc_app_resources( + self, + *, + manager: Any, + output: WebRTCOutputSpec, + spec: DemoSpec, + ) -> WebRTCAppResources: + """Return T2V browser resources for the shared WebRTC server.""" + del output, spec + return WebRTCAppResources( + model_web_resource=files("t2v").joinpath("web"), + configure_app=lambda app: _configure_t2v_webrtc_app( + app, + manager=manager, + model=self.model, + ), + preload_name="FlashDreams T2V", + ) + class T2VInputProvider: """No-control input provider for finite prompt-only generation.""" @@ -507,7 +539,12 @@ def create_t2v_spec( **defaults.runtime_options, }, ), - metadata={"output_layout": "tchw"}, + metadata={ + "output_layout": "tchw", + "webrtc_keep_connection_after_completed": True, + "webrtc_preload_name": "FlashDreams T2V", + "webrtc_supported_control_keys": ("g",), + }, ) @@ -548,6 +585,107 @@ def _validate_optional_positive_int(value: int | None, *, name: str) -> None: _validate_positive_int(value, name=name) +def _configure_t2v_webrtc_app( + app: Any, + *, + manager: Any, + model: T2VModelConfig, +) -> None: + from aiohttp import web + + selected_backend = str(model.runtime_options.get("backend", model.model_id)) + selected_label = str(model.runtime_options.get("application", model.model_id)) + + async def app_config(_: web.Request) -> web.StreamResponse: + return web.json_response( + { + "backends": [{"key": selected_backend, "label": selected_label}], + "selected_backend": selected_backend, + } + ) + + async def update_prompt(request: web.Request) -> web.StreamResponse: + payload = await request.json() + if not isinstance(payload, dict) or not isinstance(payload.get("prompt"), str): + raise web.HTTPBadRequest(reason="Expected a JSON prompt.") + duration_s = payload.get("duration_s") + if not isinstance(duration_s, int | float): + raise web.HTTPBadRequest(reason="Expected numeric duration_s.") + try: + _update_t2v_webrtc_prompt( + manager=manager, + prompt=payload["prompt"], + duration_s=float(duration_s), + ) + except (RuntimeError, ValueError) as exc: + raise web.HTTPBadRequest(reason=str(exc)) from exc + return web.json_response({"status": "ok"}) + + async def download(_: web.Request) -> web.StreamResponse: + artifact = manager.runtime.latest_artifact + if artifact is None: + raise web.HTTPNotFound(reason="No completed generation is available yet.") + video_path, scenario = artifact + if not video_path.is_file(): + raise web.HTTPNotFound(reason="Generated MP4 is no longer available.") + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + archive.write(video_path, "video.mp4") + archive.writestr( + "prompt.json", + json.dumps( + { + "prompt": scenario.prompt, + "total_blocks": scenario.total_blocks, + "fps": scenario.fps, + "width": scenario.pixel_width, + "height": scenario.pixel_height, + }, + indent=2, + ), + ) + return web.Response( + body=buffer.getvalue(), + headers={ + "Content-Disposition": "attachment; filename=flashdreams-generation.zip" + }, + content_type="application/zip", + ) + + async def playback(_: web.Request) -> web.StreamResponse: + artifact = manager.runtime.latest_artifact + if artifact is None or not artifact[0].is_file(): + raise web.HTTPNotFound(reason="No completed MP4 is available yet.") + return web.FileResponse(artifact[0]) + + app.router.add_get("/api/t2v/config", app_config) + app.router.add_post("/api/t2v/prompt", update_prompt) + app.router.add_get("/api/t2v/download", download) + app.router.add_get("/api/t2v/playback", playback) + + +def _update_t2v_webrtc_prompt( + *, + manager: Any, + prompt: str, + duration_s: float, +) -> None: + if not prompt.strip(): + raise ValueError("Prompt must be non-empty.") + if not 0 < duration_s <= 60: + raise ValueError("Duration must be greater than 0 and at most 60 seconds.") + spec = manager.shared_spec + if spec is None: + raise RuntimeError("T2V WebRTC shared session is not initialized.") + scenario = dict(spec.scenario or {}) + scenario[FIELD_PROMPT] = prompt.strip() + scenario[FIELD_TOTAL_BLOCKS] = manager.runtime.blocks_for_duration( + duration_s, + fps=_int_value(scenario[FIELD_FPS]), + ) + manager.update_shared_spec(replace(spec, scenario=scenario)) + + __all__ = [ "FIELD_FPS", "FIELD_PIXEL_HEIGHT", diff --git a/apps/t2v/web/adapter.css b/apps/t2v/web/adapter.css new file mode 100644 index 000000000..66e816f29 --- /dev/null +++ b/apps/t2v/web/adapter.css @@ -0,0 +1,3 @@ +.t2vPanel { display: grid; gap: .75rem; padding: 1rem; max-width: 32rem; } +.t2vPanel textarea { display: block; width: 100%; margin-top: .35rem; resize: vertical; } +.t2vPanel button { margin-right: .5rem; } diff --git a/apps/t2v/web/adapter.js b/apps/t2v/web/adapter.js new file mode 100644 index 000000000..7c56a622a --- /dev/null +++ b/apps/t2v/web/adapter.js @@ -0,0 +1,23 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** Model metadata only; shared WebRTC UI renders prompt/video controls. */ +export default { + modelName: "Text-to-Video", + async mount(context) { + const response = await fetch("/api/t2v/config") + if (!response.ok) return + const config = await response.json() + const selected = config.backends?.find((backend) => backend.key === config.selected_backend) + context.setModelName(selected?.label || config.selected_backend || "Text-to-Video") + }, + promptGeneration: { + endpoint: "/api/t2v/prompt", + label: "Describe the video", + placeholder: "A cinematic drone shot over snowy mountains at sunrise", + generateLabel: "Generate video", + downloadEndpoint: "/api/t2v/download", + playbackEndpoint: "/api/t2v/playback", + hideControls: true, + }, +} diff --git a/flashdreams/flashdreams/demo/__init__.py b/flashdreams/flashdreams/demo/__init__.py index 776aa7b32..a2b91df0f 100644 --- a/flashdreams/flashdreams/demo/__init__.py +++ b/flashdreams/flashdreams/demo/__init__.py @@ -7,6 +7,7 @@ DemoApplication, create_demo_application, run_application_replay, + run_application_webrtc, run_replay_application, ) from flashdreams.demo.application import ( @@ -84,5 +85,6 @@ "create_webrtc_io_handler", "input_state_from_window", "run_application_replay", + "run_application_webrtc", "run_replay_application", ] diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index f943cd7c5..37c83ebf2 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -8,6 +8,7 @@ import argparse import sys from collections.abc import Callable, Sequence +from dataclasses import dataclass, replace from typing import Any import torch @@ -18,15 +19,29 @@ configure_logging, initialize_cuda_distributed, ) +from flashdreams.runtime.demo.host import RuntimeHost from flashdreams.runtime.demo.outputs import build_output_sink from flashdreams.runtime.demo.run_modes import RunResult -from flashdreams.runtime.demo.spec import DemoAdapter, DemoSpec +from flashdreams.runtime.demo.spec import ( + DemoAdapter, + DemoSpec, + WebRTCAppResources, + WebRTCOutputSpec, +) from .application import Application, DemoAdapterApplication, IOHandler from .io import IOHandlerServer, create_replay_io_handler from .runner import Runner +@dataclass(frozen=True, slots=True) +class _PublicWebRTCConfig: + video_width: int + video_height: int + warmup_chunks: int + warmup_timeout_s: float + + class DemoApplication: """Base command application shared by model replay and WebRTC demos.""" @@ -127,7 +142,9 @@ def _run_handler(self, args: argparse.Namespace, handler: IOHandler) -> RunResul def run_replay_application(*, spec: DemoSpec, adapter: DemoAdapter) -> RunResult: """Run a finite demo spec through the public replay IO factory and runner.""" - return run_application_replay(app=DemoAdapterApplication(adapter=adapter, spec=spec)) + return run_application_replay( + app=DemoAdapterApplication(adapter=adapter, spec=spec) + ) def run_application_replay( @@ -144,6 +161,63 @@ def run_application_replay( ).run() +def run_application_webrtc( + *, app: Application, launch_args: Sequence[str] = () +) -> object: + """Serve a public ``DemoAdapterApplication`` through shared WebRTC serving.""" + del launch_args + if not isinstance(app, DemoAdapterApplication): + raise ValueError( + "Direct WebRTC application launch requires DemoAdapterApplication." + ) + output = app.spec.output + if not isinstance(output, WebRTCOutputSpec): + raise ValueError("Direct WebRTC application launch requires WebRTCOutputSpec.") + configure_logging() + config = app.spec.config + if config is None: + raise ValueError("Direct WebRTC application launch requires DemoSpec.config.") + context = initialize_cuda_distributed( + default_device=config.device or "cuda", + distributed_init_fn=distributed_init, + configure_logging_fn=configure_logging, + torch_module=torch, + dist_module=dist, + ) + web_config = replace(config, device=str(context.device)) + spec = replace( + app.spec, + input_mode="webrtc", + config=web_config, + ) + adapter = app.adapter + scenario = adapter.prepare_scenario(spec) + runtime = adapter.create_runtime(web_config) + manager = _create_application_webrtc_manager( + runtime=runtime, + output=output, + spec=spec, + scenario=scenario, + adapter=adapter, + ) + app_resources = _create_application_webrtc_resources( + adapter=adapter, + manager=manager, + output=output, + spec=spec, + ) + + from flashdreams.serving.webrtc.demo import serve_webrtc_demo + + return serve_webrtc_demo( + output=output, + model_id=spec.model_id, + session_manager=manager, + app_resources=app_resources, + world_rank=context.world_rank, + ) + + def create_demo_application( *, parse_args: Callable[[list[str] | None], argparse.Namespace], @@ -160,6 +234,81 @@ def create_demo_application( ) +def _create_application_webrtc_manager( + *, + runtime: Any, + output: WebRTCOutputSpec, + spec: DemoSpec, + scenario: Any, + adapter: DemoAdapter, +) -> Any: + from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager + + return BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=_PublicWebRTCConfig( + video_width=output.video_width, + video_height=output.video_height, + warmup_chunks=output.warmup_chunks, + warmup_timeout_s=output.warmup_timeout_s, + ), + fps=output.fps, + identity=str(getattr(adapter, "model_id", spec.model_id)), + warmup_label=output.preload_name + or _metadata_str(spec, "webrtc_preload_name", default="WebRTC"), + supported_control_keys=_metadata_string_set( + spec, "webrtc_supported_control_keys" + ), + client_liveness_timeout_s=output.client_liveness_timeout_s, + shared_host=RuntimeHost(runtime), + shared_adapter=adapter, + shared_spec=spec, + shared_scenario=scenario, + keep_connection_after_completed=bool( + spec.metadata.get("webrtc_keep_connection_after_completed", False) + ), + ) + + +def _create_application_webrtc_resources( + *, + adapter: DemoAdapter, + manager: Any, + output: WebRTCOutputSpec, + spec: DemoSpec, +) -> WebRTCAppResources: + factory = getattr(adapter, "create_webrtc_app_resources", None) + if callable(factory): + resources = factory(manager=manager, output=output, spec=spec) + if not isinstance(resources, WebRTCAppResources): + raise TypeError( + "Demo adapter create_webrtc_app_resources(...) must return " + f"WebRTCAppResources, got {type(resources).__name__}." + ) + return resources + resources = spec.metadata.get("webrtc_app_resources") + if isinstance(resources, WebRTCAppResources): + return resources + return WebRTCAppResources(preload_name=output.preload_name or spec.model_id) + + +def _metadata_str(spec: DemoSpec, name: str, *, default: str) -> str: + value = spec.metadata.get(name, default) + return str(value) + + +def _metadata_string_set(spec: DemoSpec, name: str) -> frozenset[str] | None: + value = spec.metadata.get(name) + if value is None: + return None + if isinstance(value, str): + return frozenset({value}) + try: + return frozenset(str(item) for item in value) + except TypeError as exc: + raise TypeError(f"DemoSpec.metadata[{name!r}] must be iterable.") from exc + + def _raise_for_failed_result(result: RunResult) -> None: if result.status in {"completed", "skipped"}: return @@ -173,6 +322,7 @@ def _raise_for_failed_result(result: RunResult) -> None: __all__ = [ "DemoApplication", "create_demo_application", + "run_application_webrtc", "run_application_replay", "run_replay_application", ] diff --git a/flashdreams/flashdreams/runtime/video_runner.py b/flashdreams/flashdreams/runtime/video_runner.py new file mode 100644 index 000000000..444b7c022 --- /dev/null +++ b/flashdreams/flashdreams/runtime/video_runner.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared prompt-conditioned video runners for ``flashdreams-run`` slugs. + +Every model that generates video from a prompt drives the same rollout: build +a cache, call ``generate``/``finalize`` per autoregressive index, stream chunks +into an MP4, then write per-step stats. Integrations declare their defaults as +a :class:`VideoRunnerConfig` subclass in their ``config.py`` and point +``_target`` at one of the runners here instead of restating the rollout. + +Pipelines are consumed structurally rather than by type: a pipeline must expose +``initialize_cache``, ``generate(autoregressive_index=..., cache=...)``, +``finalize(autoregressive_index=..., cache=...)``, and a ``decoder`` that is a +:class:`~flashdreams.infra.decoder.StreamingVideoDecoder`. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, ClassVar + +import torch +from loguru import logger + +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + load_first_frame_tensor, + read_image_rgb, + resolve_input_path, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, +) +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.video_output import Mp4VideoOutputTarget + +__all__ = [ + "ImageConditionedVideoRunnerConfig", + "StreamingVideoRunner", + "StreamingVideoRunnerConfig", + "VideoRunner", + "VideoRunnerConfig", + "image_cache_dir", +] + + +def image_cache_dir(subdir: str) -> Path: + """Return the user-writable cache for on-the-fly image downloads.""" + root = os.path.expanduser( + os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams") + ) + return Path(root) / subdir + + +@dataclass(kw_only=True) +class VideoRunnerConfig(RunnerConfig): + """Base config for a prompt-conditioned single-step video runner.""" + + _target: type["VideoRunner"] = field(default_factory=lambda: VideoRunner) + + prompt: str | Path = "" + """Either an inline text prompt (--prompt "...") or a path to a + txt file whose first line is read as the prompt (--prompt prompt.txt).""" + + pixel_height: int = 480 + """Output video pixel height.""" + + pixel_width: int = 832 + """Output video pixel width.""" + + fps: int = 16 + """Output video frame rate.""" + + postprocess_output_layout: VideoTensorLayout | None = "tchw" + """Pipeline output layout for streaming post-processing.""" + + +@dataclass(kw_only=True) +class StreamingVideoRunnerConfig(VideoRunnerConfig): + """Config for autoregressive models that roll out many chunks.""" + + _target: type["StreamingVideoRunner"] = field( + default_factory=lambda: StreamingVideoRunner + ) + + total_blocks: int = 60 + """Number of autoregressive chunks to generate before terminating the rollout.""" + + +@dataclass(kw_only=True) +class ImageConditionedVideoRunnerConfig: + """Mixin adding the first-frame image that I2V variants need at runtime. + + Inherit it alongside :class:`VideoRunnerConfig` or + :class:`StreamingVideoRunnerConfig`; the runners pick these fields up when + they are present so that T2V slugs keep an image-free CLI surface. + """ + + image_path: str | Path = "" + """First-frame RGB image. Either a local path or an HTTP(S) URL.""" + + image_cache_subdir: ClassVar[str] = "video" + """Subdirectory of the FlashDreams cache for downloaded images. + + A per-model constant rather than a field, so it stays off the CLI. + """ + + +class VideoRunner(Runner[VideoRunnerConfig, Any]): + """Prompt-conditioned video runner that generates one chunk.""" + + config: VideoRunnerConfig + + def _step_count(self) -> int: + return 1 + + def _resolve_prompt(self) -> str: + """Resolve ``config.prompt``. + + A Path reads its first non-empty line, a str is used as-is. + """ + return resolve_prompt_value(self.config.prompt) + + def _latent_dimensions(self) -> tuple[int, int]: + """Return the latent height and width for the configured pixel size.""" + config = self.config + decoder = self.pipeline.decoder + if not isinstance(decoder, StreamingVideoDecoder): + raise TypeError( + f"[{config.runner_name}] requires a StreamingVideoDecoder, " + f"got {type(decoder).__name__}." + ) + ratio = decoder.spatial_compression_ratio + if config.pixel_height % ratio or config.pixel_width % ratio: + raise ValueError( + f"[{config.runner_name}] pixel_height={config.pixel_height} and " + f"pixel_width={config.pixel_width} must both divide {ratio}." + ) + return config.pixel_height // ratio, config.pixel_width // ratio + + def _conditioning_image(self) -> torch.Tensor | None: + """Load the first-frame image when this model is image-conditioned. + + Read structurally so that T2V configs, which do not declare the mixin + fields, keep an image-free CLI surface. + """ + image_path = getattr(self.config, "image_path", "") + if not image_path: + return None + config = self.config + # Load + resize the first frame, then convert to [-1, 1] bf16 in shape + # [T=1, C, H, W]. Pin to the pipeline's actual device so non-default + # ``--device`` selections and the torchrun cuda:LOCAL_RANK override + # both work. + return load_first_frame_tensor( + resolve_input_path( + image_path, + cache_dir=image_cache_dir( + getattr(config, "image_cache_subdir", "video") + ), + validator=read_image_rgb, + ), + pixel_height=config.pixel_height, + pixel_width=config.pixel_width, + device=self.pipeline.device, + dtype=torch.bfloat16, + ) + + def _initialize_cache(self) -> Any: + """Initialize the rollout cache for either T2V or I2V conditioning.""" + prompt = self._resolve_prompt() + latent_height, latent_width = self._latent_dimensions() + image = self._conditioning_image() + if image is not None: + return self.pipeline.initialize_cache(text=[prompt], image=image) + return self.pipeline.initialize_cache( + text=[prompt], image=None, height=latent_height, width=latent_width + ) + + def run(self) -> None: + """Drive the rollout and write the video plus per-step stats.""" + config = self.config + cache = self._initialize_cache() + output_stream = self.create_video_output_stream(fps=config.fps) + output_target = Mp4VideoOutputTarget( + output_path=runner_artifact_path( + config.output_dir, config.runner_name, "mp4" + ), + fps=config.fps, + output_layout=output_stream.output_layout, + enabled=self.is_rank_zero, + ) + output_target.open() + for index in range(self._step_count()): + chunk = self.pipeline.generate(autoregressive_index=index, cache=cache) + stats = self.pipeline.finalize(autoregressive_index=index, cache=cache) + output_target.write( + output_stream.process(chunk, autoregressive_index=index, metrics=stats) + ) + tail = output_stream.finish() + if tail is not None: + output_target.write(tail) + artifacts = output_target.close() + if artifacts: + self._log_artifact(artifacts[0]) + + def _log_artifact(self, video_artifact: OutputArtifact) -> None: + """Log the written video and persist per-step stats when present.""" + config = self.config + logger.info( + f"[{config.runner_name}] wrote video " + f"{video_artifact.metadata['shape']} " + f"-> {Path(video_artifact.uri).resolve()}" + ) + stats_history = video_artifact.metadata["stats_history"] + if not stats_history: + return + stats_path = write_runner_stats( + config.output_dir, config.runner_name, list(stats_history) + ) + logger.info( + f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" + ) + + +class StreamingVideoRunner(VideoRunner): + """Prompt-conditioned video runner that rolls out ``total_blocks`` chunks.""" + + config: StreamingVideoRunnerConfig + + def _step_count(self) -> int: + return self.config.total_blocks diff --git a/flashdreams/flashdreams/scripts/cli.py b/flashdreams/flashdreams/scripts/cli.py index afb246e80..fd794b917 100644 --- a/flashdreams/flashdreams/scripts/cli.py +++ b/flashdreams/flashdreams/scripts/cli.py @@ -57,10 +57,17 @@ Application, DemoAdapterApplication, run_application_replay, + run_application_webrtc, ) from flashdreams.infra.runner import RunnerConfig from flashdreams.plugins import discover_applications -from flashdreams.runtime.demo import DemoSpec, Mp4OutputSpec, NullOutputSpec +from flashdreams.runtime.demo import ( + DemoAdapter, + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) from flashdreams.serving.launch import ( LaunchMode, LaunchOptions, @@ -386,18 +393,18 @@ def _entrypoint_application( if "--help" in raw_args or "-h" in raw_args: _print_application_help(application_name, application) raise SystemExit(0) - mode, launch_args, no_instantiate, launch_overrides = ( - _prepare_application_cli_args(raw_args, application_name=application_name) + mode, launch_args, no_instantiate, launch_overrides = _prepare_application_cli_args( + raw_args, application_name=application_name ) - if mode not in {"mp4", "null"}: + if mode not in {"mp4", "null", "webrtc"}: raise ValueError( - f"Application {application_name!r} currently supports direct finite " - "launch modes 'mp4' and 'null'. Use a compatibility runner for " + f"Application {application_name!r} currently supports direct launch " + "modes 'mp4', 'null', and 'webrtc'. Use a compatibility runner for " f"{mode!r}." ) scenario = dict(launch_overrides.scenario) output = dict(launch_overrides.output) - configured = _configure_application_replay( + configured = _configure_application_launch( application=application, application_name=application_name, mode=mode, @@ -413,9 +420,14 @@ def _entrypoint_application( print(f"Output settings: {output}") if no_instantiate: return - _handle_launch_result( - run_application_replay(app=configured, launch_args=launch_args) - ) + if mode == "webrtc": + _handle_launch_result( + run_application_webrtc(app=configured, launch_args=launch_args) + ) + else: + _handle_launch_result( + run_application_replay(app=configured, launch_args=launch_args) + ) def _prepare_application_cli_args( @@ -430,6 +442,18 @@ def _prepare_application_cli_args( "--manifest is not supported for direct application launches yet; " "use --scenario.KEY and --output.KEY overrides." ) + normalized, host = _pop_option(normalized, "--host") + normalized, port = _pop_option(normalized, "--port") + if host is not None or port is not None: + output_overrides = dict(launch_overrides.output) + if host is not None: + output_overrides["host"] = host + if port is not None: + output_overrides["port"] = port + launch_overrides = dataclasses.replace( + launch_overrides, + output=output_overrides, + ) normalized = _hoist_global_options(normalized) no_instantiate = False remaining: list[str] = [] @@ -444,12 +468,6 @@ def _prepare_application_cli_args( raise ValueError( "--prefer-sw-encoder is only supported by WebRTC runner launches." ) - if item in {"--host", "--port"}: - raise ValueError(f"{item} is only supported by WebRTC runner launches.") - if any(item.startswith(option + "=") for option in ("--host", "--port")): - raise ValueError( - f"{item.split('=', 1)[0]} is only supported by WebRTC runner launches." - ) remaining.append(item) index += 1 @@ -466,7 +484,7 @@ def _prepare_application_cli_args( return raw_mode, tuple(remaining), no_instantiate, launch_overrides -def _configure_application_replay( +def _configure_application_launch( *, application: Application, application_name: str, @@ -490,17 +508,28 @@ def _configure_application_replay( scenario=scenario, output_overrides=output_overrides, ) + adapter = _application_adapter_for_output(application.adapter, output) return DemoAdapterApplication( - adapter=application.adapter, + adapter=adapter, spec=dataclasses.replace( application.spec, - input_mode="replay", + input_mode="webrtc" if mode == "webrtc" else "replay", scenario=scenario, output=output, ), ) +def _application_adapter_for_output( + adapter: DemoAdapter, + output: object, +) -> DemoAdapter: + configure_for_output = getattr(adapter, "configure_for_output", None) + if not callable(configure_for_output): + return adapter + return cast(DemoAdapter, configure_for_output(output)) + + def _application_scenario( spec: DemoSpec, overrides: Mapping[str, object], @@ -524,16 +553,20 @@ def _application_output_spec( spec: DemoSpec, scenario: object, output_overrides: Mapping[str, object], -) -> Mp4OutputSpec | NullOutputSpec: +) -> Mp4OutputSpec | NullOutputSpec | WebRTCOutputSpec: if mode == "null": _reject_unknown_output_keys(output_overrides, allowed={"store_results"}) return NullOutputSpec( store_results=bool(output_overrides.get("store_results", False)) ) - if mode != "mp4": - raise ValueError( - f"Direct application launch mode {mode!r} is not implemented." + if mode == "webrtc": + return _application_webrtc_output_spec( + spec=spec, + scenario=scenario, + output_overrides=output_overrides, ) + if mode != "mp4": + raise ValueError(f"Direct application launch mode {mode!r} is not implemented.") _reject_unknown_output_keys( output_overrides, allowed={"fps", "layout", "move_to_cpu", "output", "output_layout", "path"}, @@ -575,6 +608,150 @@ def _application_output_spec( ) +def _application_webrtc_output_spec( + *, + spec: DemoSpec, + scenario: object, + output_overrides: Mapping[str, object], +) -> WebRTCOutputSpec: + _reject_unknown_output_keys( + output_overrides, + allowed={ + "client_liveness_timeout_s", + "fps", + "host", + "port", + "preload_name", + "request_session_path", + "video_height", + "video_width", + "warmup_chunks", + "warmup_timeout_s", + "web_dir", + }, + ) + current_output = spec.output + return WebRTCOutputSpec( + host=str( + _webrtc_output_value( + output_overrides, + current_output, + "host", + default="127.0.0.1", + ) + ), + port=_positive_int( + _webrtc_output_value( + output_overrides, + current_output, + "port", + default=8080, + ), + name="port", + ), + fps=_positive_int( + _webrtc_output_value( + output_overrides, + current_output, + "fps", + default=( + _scenario_field(scenario, "fps") or spec.metadata.get("fps") or 30 + ), + ), + name="fps", + ), + video_width=_positive_int( + _webrtc_output_value( + output_overrides, + current_output, + "video_width", + default=( + _scenario_field(scenario, "pixel_width") + or spec.metadata.get("video_width") + or 1280 + ), + ), + name="video_width", + ), + video_height=_positive_int( + _webrtc_output_value( + output_overrides, + current_output, + "video_height", + default=( + _scenario_field(scenario, "pixel_height") + or spec.metadata.get("video_height") + or 720 + ), + ), + name="video_height", + ), + warmup_chunks=_non_negative_int( + _webrtc_output_value( + output_overrides, + current_output, + "warmup_chunks", + default=0, + ), + name="warmup_chunks", + ), + warmup_timeout_s=float( + _positive_number( + _webrtc_output_value( + output_overrides, + current_output, + "warmup_timeout_s", + default=30.0, + ), + name="warmup_timeout_s", + ) + ), + client_liveness_timeout_s=float( + _positive_number( + _webrtc_output_value( + output_overrides, + current_output, + "client_liveness_timeout_s", + default=30.0, + ), + name="client_liveness_timeout_s", + ) + ), + web_dir=_optional_path( + _webrtc_output_value(output_overrides, current_output, "web_dir") + ), + request_session_path=str( + _webrtc_output_value( + output_overrides, + current_output, + "request_session_path", + default="/request_session", + ) + ), + preload_name=_optional_str( + _webrtc_output_value( + output_overrides, + current_output, + "preload_name", + ) + ), + ) + + +def _webrtc_output_value( + output_overrides: Mapping[str, object], + current_output: object, + name: str, + *, + default: object | None = None, +) -> object | None: + if name in output_overrides: + return output_overrides[name] + if isinstance(current_output, WebRTCOutputSpec): + return getattr(current_output, name) + return default + + def _scenario_field(scenario: object, name: str) -> object | None: if isinstance(scenario, Mapping): return scenario.get(name) @@ -595,6 +772,39 @@ def _positive_number(value: object, *, name: str) -> int | float: return number +def _positive_int(value: object, *, name: str) -> int: + number = _positive_number(value, name=name) + return int(number) + + +def _non_negative_int(value: object, *, name: str) -> int: + if isinstance(value, bool): + raise TypeError(f"{name} must be an integer, not bool.") + if isinstance(value, int): + number = value + elif isinstance(value, float): + number = int(value) + elif isinstance(value, str): + number = int(value) + else: + raise TypeError(f"{name} must be an integer, got {type(value).__name__}.") + if number < 0: + raise ValueError(f"{name} must be >= 0.") + return number + + +def _optional_path(value: object | None) -> Path | None: + if value is None: + return None + return Path(str(value)) + + +def _optional_str(value: object | None) -> str | None: + if value is None: + return None + return str(value) + + def _reject_unknown_output_keys( output_overrides: Mapping[str, object], *, @@ -609,10 +819,10 @@ def _print_application_help(application_name: str, application: Application) -> modes = ("null",) if isinstance(application, DemoAdapterApplication): supported = set(application.adapter.supported_output_modes()) - modes = tuple(mode for mode in ("mp4", "null") if mode in supported) + modes = tuple(mode for mode in ("mp4", "null", "webrtc") if mode in supported) print(f"Usage: flashdreams-run {application_name} [options]") print(f"Available direct application modes: {', '.join(modes)}") - print("Use --scenario.KEY VALUE and --output.KEY VALUE for finite-mode settings.") + print("Use --scenario.KEY VALUE and --output.KEY VALUE for mode settings.") def _prepare_cli_args( diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 149f83d4f..44190e592 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -728,6 +728,19 @@ def runtime(self) -> _RuntimeT: """Model runtime driven by this transport manager.""" return self._runtime + @property + def shared_spec(self) -> DemoSpec | None: + """Current shared spec used to create new WebRTC sessions.""" + return self._shared_spec + + def update_shared_spec(self, spec: DemoSpec) -> None: + """Replace the shared spec and rebuild the prepared scenario.""" + adapter = self._shared_adapter + if adapter is None: + raise RuntimeError("This WebRTC manager has no shared demo adapter.") + self._shared_spec = spec + self._shared_scenario = adapter.prepare_scenario(spec) + def set_pending_session_input(self, session_input: Any) -> None: """Store validated model input for the next session.""" if self.has_active_session(): diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 5ce4f0523..0e54e3c86 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -39,6 +39,7 @@ create_webrtc_io_handler, input_state_from_window, run_application_replay, + run_application_webrtc, ) from flashdreams.runtime import ( CanonicalInputSchema, @@ -95,6 +96,7 @@ def test_public_demo_contracts_are_importable() -> None: assert IOHandlerServer.__name__ == "IOHandlerServer" assert FrameOutputSink.__name__ == "FrameOutputSink" assert run_application_replay.__name__ == "run_application_replay" + assert run_application_webrtc.__name__ == "run_application_webrtc" assert create_demo_application.__name__ == "create_demo_application" assert create_replay_io_handler.__name__ == "create_replay_io_handler" assert create_native_window_io_handler.__name__ == ( diff --git a/flashdreams/tests/test_launch_manifest.py b/flashdreams/tests/test_launch_manifest.py index bd572c3cb..bbcf6537d 100644 --- a/flashdreams/tests/test_launch_manifest.py +++ b/flashdreams/tests/test_launch_manifest.py @@ -17,6 +17,7 @@ InferenceConfig, InferenceInput, InferenceInputSchema, + InferenceRuntime, InputMapping, ) from flashdreams.runtime.demo import ( @@ -25,6 +26,7 @@ NullOutputSpec, PreparedScenario, RunResult, + WebRTCOutputSpec, ) from flashdreams.scripts import cli from flashdreams.serving.launch import ResolvedLaunch, resolve_launch @@ -148,6 +150,49 @@ def test_entrypoint_application_null_rejects_output_path( cli.entrypoint(["test-app", "null", "--output.path", "unexpected.mp4"]) +def test_entrypoint_launches_application_slug_webrtc_directly( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: list[tuple[DemoAdapterApplication, tuple[str, ...]]] = [] + + def fake_run_application_webrtc( + *, app: object, launch_args: tuple[str, ...] = () + ) -> RunResult: + assert isinstance(app, DemoAdapterApplication) + captured.append((app, launch_args)) + return RunResult(status="completed") + + monkeypatch.setattr(cli, "all_runners", lambda: {}) + monkeypatch.setattr( + cli, + "discover_applications", + lambda: {"test-app": _application()}, + ) + monkeypatch.setattr(cli, "run_application_webrtc", fake_run_application_webrtc) + + cli.entrypoint( + [ + "test-app", + "webrtc", + "--host", + "0.0.0.0", + "--port", + "8089", + "--scenario.fps", + "12", + ] + ) + + configured, launch_args = captured[0] + assert launch_args == () + assert configured.spec.input_mode == "webrtc" + output = configured.spec.output + assert isinstance(output, WebRTCOutputSpec) + assert output.host == "0.0.0.0" + assert output.port == 8089 + assert output.fps == 12 + + def test_launch_manifest_does_not_guess_configs_directory( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -540,10 +585,10 @@ class _CliApplicationAdapter: canonical_input_schema = CanonicalInputSchema() def supported_input_modes(self) -> tuple[str, ...]: - return ("replay",) + return ("replay", "webrtc") def supported_output_modes(self) -> tuple[str, ...]: - return ("mp4", "null") + return ("mp4", "null", "webrtc") def default_input_mapping(self) -> InputMapping: return IdentityInputMapping() @@ -555,6 +600,6 @@ def prepare_scenario(self, spec: DemoSpec) -> PreparedScenario: del spec return PreparedScenario(initial_inputs=InferenceInput()) - def create_runtime(self, config: InferenceConfig) -> object: + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: del config raise AssertionError("direct app CLI test should not instantiate the runtime") diff --git a/integrations/causal_forcing/causal_forcing/config.py b/integrations/causal_forcing/causal_forcing/config.py index 98242d8d9..b3fa86489 100644 --- a/integrations/causal_forcing/causal_forcing/config.py +++ b/integrations/causal_forcing/causal_forcing/config.py @@ -17,14 +17,12 @@ from __future__ import annotations +from dataclasses import dataclass +from pathlib import Path from typing import Any, cast from torch import Tensor -from causal_forcing.runner import ( - CausalForcingI2VRunnerConfig, - CausalForcingT2VRunnerConfig, -) from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig @@ -37,10 +35,49 @@ WanVAEDecoderConfig, WanVAEEncoderConfig, ) +from flashdreams.runtime.video_runner import ( + ImageConditionedVideoRunnerConfig, + StreamingVideoRunnerConfig, +) CHECKPOINT_PATH_CHUNKWISE = "https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/chunkwise/causal_forcing.pt" CHECKPOINT_PATH_FRAMEWISE = "https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/framewise/causal_forcing.pt" +DEFAULT_T2V_PROMPT = ( + "A cinematic closeup and detailed portrait of a reindeer standing in a " + "snowy forest at sunset. The lighting is gorgeous and soft, with a golden " + "backlight creating a warm and dreamy effect. Soft bokeh and lens flares " + "add a magical touch, enhancing the cinematic quality of the image. The " + "reindeer has a gentle expression, its fur glistening in the fading light. " + "The background features a serene snowy landscape with tall trees " + "silhouetted against the orange and pink hues of the setting sun. The " + "color grade is rich and magical, capturing the essence of a winter " + "wonderland at twilight. A close-up shot from a slightly elevated angle." +) + +DEFAULT_I2V_IMAGE_URL = "https://raw.githubusercontent.com/thu-ml/Causal-Forcing/refs/heads/main/prompts/i2v/26-15/000001.png" + + +@dataclass(kw_only=True) +class CausalForcingT2VRunnerConfig(StreamingVideoRunnerConfig): + """Runner config for the Causal-Forcing T2V variants.""" + + prompt: str | Path = DEFAULT_T2V_PROMPT + total_blocks: int = 60 + pixel_height: int = 480 + pixel_width: int = 832 + fps: int = 16 + + +@dataclass(kw_only=True) +class CausalForcingI2VRunnerConfig( + ImageConditionedVideoRunnerConfig, CausalForcingT2VRunnerConfig +): + """Runner config for the Causal-Forcing I2V variants.""" + + image_path: str | Path = DEFAULT_I2V_IMAGE_URL + image_cache_subdir = "causal_forcing" + def state_dict_transform(state_dict: dict[str, Any]) -> dict[str, Tensor]: """Strip Causal-Forcing wrapper prefixes from the checkpoint state-dict. diff --git a/integrations/causal_forcing/causal_forcing/runner.py b/integrations/causal_forcing/causal_forcing/runner.py deleted file mode 100644 index 574ec0561..000000000 --- a/integrations/causal_forcing/causal_forcing/runner.py +++ /dev/null @@ -1,256 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Causal-Forcing Wan 2.1 streaming runner classes (T2V and I2V).""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from pathlib import Path - -import torch -from loguru import logger - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - load_first_frame_tensor, - read_image_rgb, - resolve_input_path, - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) -from flashdreams.recipes.wan import ( - WanInferencePipeline, - WanInferencePipelineCache, -) -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - -__all__ = [ - "CausalForcingI2VRunnerConfig", - "CausalForcingI2VRunner", - "CausalForcingT2VRunnerConfig", - "CausalForcingT2VRunner", -] - - -DEFAULT_T2V_PROMPT = ( - "A cinematic closeup and detailed portrait of a reindeer standing in a " - "snowy forest at sunset. The lighting is gorgeous and soft, with a golden " - "backlight creating a warm and dreamy effect. Soft bokeh and lens flares " - "add a magical touch, enhancing the cinematic quality of the image. The " - "reindeer has a gentle expression, its fur glistening in the fading light. " - "The background features a serene snowy landscape with tall trees " - "silhouetted against the orange and pink hues of the setting sun. The " - "color grade is rich and magical, capturing the essence of a winter " - "wonderland at twilight. A close-up shot from a slightly elevated angle." -) - - -DEFAULT_I2V_IMAGE_URL = "https://raw.githubusercontent.com/thu-ml/Causal-Forcing/refs/heads/main/prompts/i2v/26-15/000001.png" - -IMAGE_CACHE_DIR = ( - Path(os.path.expanduser(os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams"))) - / "self_forcing" -) -"""User-writable cache for on-the-fly I2V first-frame downloads.""" - - -@dataclass(kw_only=True) -class CausalForcingT2VRunnerConfig(RunnerConfig): - """Runner config for the Causal-Forcing T2V variants. - - Also serves as the base for :class:`CausalForcingI2VRunnerConfig` - (I2V is T2V plus an ``image_path``). - """ - - _target: type["CausalForcingT2VRunner"] = field( - default_factory=lambda: CausalForcingT2VRunner - ) - - prompt: str | Path = DEFAULT_T2V_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt).""" - - total_blocks: int = 60 - """Number of autoregressive chunks to generate before terminating the rollout.""" - - pixel_height: int = 480 - """Output video pixel height.""" - - pixel_width: int = 832 - """Output video pixel width.""" - - fps: int = 16 - """Output video frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Pipeline output layout for streaming post-processing.""" - - -@dataclass(kw_only=True) -class CausalForcingI2VRunnerConfig(CausalForcingT2VRunnerConfig): - """Runner config for the Causal-Forcing I2V variants. - - Inherits all T2V fields (prompt, total_blocks, pixel_*, fps) and - adds the first-frame image path that I2V needs at runtime. - """ - - _target: type["CausalForcingI2VRunner"] = field( - default_factory=lambda: CausalForcingI2VRunner - ) - - image_path: str | Path = DEFAULT_I2V_IMAGE_URL - """First-frame RGB image. Either a local path or an HTTP(S) URL.""" - - -class CausalForcingT2VRunner( - Runner[CausalForcingT2VRunnerConfig, WanInferencePipeline] -): - """Causal-Forcing Wan 2.1 streaming T2V driver. - - Also serves as the base for :class:`CausalForcingI2VRunner` (I2V - only overrides :meth:`_initialize_cache` to load the first frame; - everything else, including :meth:`run`, is reused). - """ - - config: CausalForcingT2VRunnerConfig - - def _resolve_prompt(self) -> str: - """Resolve config.prompt. - - A Path reads its first non-empty line, a str is used as-is. - """ - return resolve_prompt_value(self.config.prompt) - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache for T2V.""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - latent_h = config.pixel_height // sp - latent_w = config.pixel_width // sp - - return self.pipeline.initialize_cache( - text=[prompt], image=None, height=latent_h, width=latent_w - ) - - def run(self) -> None: - """Drive the autoregressive rollout and write outputs.""" - config = self.config - - # Initialize the autoregressive cache. - cache = self._initialize_cache() - - # Generate the autoregressive chunks. - output_stream = self.create_video_output_stream(fps=config.fps) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - for i in range(config.total_blocks): - video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_target.write( - output_stream.process( - video_chunk, - autoregressive_index=i, - metrics=stats, - ) - ) - - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - video_artifact = artifacts[0] - video_path = Path(video_artifact.uri) - - logger.info( - f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " - f"-> {video_path.resolve()}" - ) - - # Write the perf stats. - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) - - -class CausalForcingI2VRunner(CausalForcingT2VRunner): - """Causal-Forcing Wan 2.1 streaming I2V driver (mask-injection first frame). - - Inherits :meth:`run` and :meth:`_resolve_prompt` from - :class:`CausalForcingT2VRunner`; only :meth:`_initialize_cache` - differs (loads + encodes the first frame). - """ - - config: CausalForcingI2VRunnerConfig - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache for I2V (loads first frame).""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - - # Load + resize the first frame, then convert to [-1, 1] bf16 - # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the - # pipeline's actual device so non-default ``--device`` selections - # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - image = load_first_frame_tensor( - resolve_input_path( - config.image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=read_image_rgb, - ), - pixel_height=config.pixel_height, - pixel_width=config.pixel_width, - device=self.pipeline.device, - dtype=torch.bfloat16, - ) - - return self.pipeline.initialize_cache(text=[prompt], image=image) diff --git a/integrations/cosmos_predict2/cosmos_predict2/config.py b/integrations/cosmos_predict2/cosmos_predict2/config.py index e469708a8..cc59dc65b 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/config.py +++ b/integrations/cosmos_predict2/cosmos_predict2/config.py @@ -17,7 +17,52 @@ from __future__ import annotations -from cosmos_predict2.runner import Cosmos2I2VRunnerConfig, Cosmos2T2VRunnerConfig +from dataclasses import dataclass +from pathlib import Path + +from flashdreams.runtime.video_runner import ( + ImageConditionedVideoRunnerConfig, + VideoRunnerConfig, +) + +DEFAULT_PROMPT = ( + "A high-definition video captures the precision of robotic welding in an industrial setting. " + "The first frame showcases a robotic arm, equipped with a welding torch, positioned over a " + "large metal structure. The welding process is in full swing, with bright sparks and intense " + "light illuminating the scene, creating a vivid display of blue and white hues. A significant " + "amount of smoke billows around the welding area, partially obscuring the view but emphasizing " + "the heat and activity. The background reveals parts of the workshop environment, including a " + "ventilation system and various pieces of machinery, indicating a busy and functional industrial " + "workspace. As the video progresses, the robotic arm maintains its steady position, continuing " + "the welding process and moving to its left. The welding torch consistently emits sparks and light, " + "and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath " + "the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, " + "with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of " + "the welding operation." +) +"""Default demo prompt used when no ``--prompt`` is supplied.""" + +DEFAULT_I2V_IMAGE_URL = "https://media.githubusercontent.com/media/nvidia-cosmos/cosmos-predict2.5/refs/heads/main/assets/base/robot_welding.jpg" + + +@dataclass(kw_only=True) +class Cosmos2T2VRunnerConfig(VideoRunnerConfig): + """Runner config for the Cosmos-Predict2 T2V variant.""" + + prompt: str | Path = DEFAULT_PROMPT + pixel_height: int = 720 + pixel_width: int = 1280 + fps: int = 16 + + +@dataclass(kw_only=True) +class Cosmos2I2VRunnerConfig(ImageConditionedVideoRunnerConfig, Cosmos2T2VRunnerConfig): + """Runner config for the Cosmos-Predict2 I2V variant.""" + + image_path: str | Path = DEFAULT_I2V_IMAGE_URL + image_cache_subdir = "cosmos_predict2" + + from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler import ( FlowMatchUniPCSchedulerConfig, diff --git a/integrations/cosmos_predict2/cosmos_predict2/runner.py b/integrations/cosmos_predict2/cosmos_predict2/runner.py deleted file mode 100644 index d37bbad69..000000000 --- a/integrations/cosmos_predict2/cosmos_predict2/runner.py +++ /dev/null @@ -1,225 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Non-streaming Cosmos-Predict2 T2V runner.""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from pathlib import Path - -import torch -from loguru import logger - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - load_first_frame_tensor, - read_image_rgb, - resolve_input_path, - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) -from flashdreams.recipes.cosmos.pipeline import ( - CosmosInferencePipeline, - CosmosInferencePipelineCache, -) -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - -__all__ = [ - "Cosmos2I2VRunner", - "Cosmos2I2VRunnerConfig", - "Cosmos2T2VRunner", - "Cosmos2T2VRunnerConfig", -] - -DEFAULT_PROMPT = ( - "A high-definition video captures the precision of robotic welding in an industrial setting. " - "The first frame showcases a robotic arm, equipped with a welding torch, positioned over a " - "large metal structure. The welding process is in full swing, with bright sparks and intense " - "light illuminating the scene, creating a vivid display of blue and white hues. A significant " - "amount of smoke billows around the welding area, partially obscuring the view but emphasizing " - "the heat and activity. The background reveals parts of the workshop environment, including a " - "ventilation system and various pieces of machinery, indicating a busy and functional industrial " - "workspace. As the video progresses, the robotic arm maintains its steady position, continuing " - "the welding process and moving to its left. The welding torch consistently emits sparks and light, " - "and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath " - "the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, " - "with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of " - "the welding operation." -) -"""Default demo prompt used when no ``--prompt`` is supplied.""" - - -DEFAULT_I2V_IMAGE_URL = "https://media.githubusercontent.com/media/nvidia-cosmos/cosmos-predict2.5/refs/heads/main/assets/base/robot_welding.jpg" - -IMAGE_CACHE_DIR = ( - Path(os.path.expanduser(os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams"))) - / "cosmos_predict2" -) -"""User-writable cache for on-the-fly I2V first-frame downloads.""" - - -@dataclass(kw_only=True) -class Cosmos2T2VRunnerConfig(RunnerConfig): - """Runner config for the Cosmos-Predict2 T2V variant.""" - - _target: type["Cosmos2T2VRunner"] = field(default_factory=lambda: Cosmos2T2VRunner) - - prompt: str | Path = DEFAULT_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt).""" - - pixel_height: int = 720 - """Output video pixel height.""" - - pixel_width: int = 1280 - """Output video pixel width.""" - - fps: int = 16 - """Output video frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Pipeline output layout for streaming post-processing.""" - - -class Cosmos2T2VRunner(Runner[Cosmos2T2VRunnerConfig, CosmosInferencePipeline]): - """Cosmos-Predict2 non-streaming T2V driver.""" - - config: Cosmos2T2VRunnerConfig - - def _resolve_prompt(self) -> str: - """Resolve config.prompt. - - A Path reads its first non-empty line, a str is used as-is. - """ - return resolve_prompt_value(self.config.prompt) - - def _initialize_cache(self) -> CosmosInferencePipelineCache: - """Initialize the autoregressive cache for T2V.""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - latent_h = config.pixel_height // sp - latent_w = config.pixel_width // sp - - return self.pipeline.initialize_cache( - text=[prompt], image=None, height=latent_h, width=latent_w - ) - - def run(self) -> None: - """Drive the single-step rollout and write outputs.""" - config = self.config - - cache = self._initialize_cache() - - output_stream = self.create_video_output_stream(fps=config.fps) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - generated = self.pipeline.generate(autoregressive_index=0, cache=cache) - stats = self.pipeline.finalize(autoregressive_index=0, cache=cache) - output_target.write( - output_stream.process(generated, autoregressive_index=0, metrics=stats) - ) - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - video_artifact = artifacts[0] - video_path = Path(video_artifact.uri) - - logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " - f"-> {video_path.resolve()}" - ) - - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - f"[{config.runner_name}] wrote per-AR-step stats " - f"-> {stats_path.resolve()}" - ) - - -@dataclass(kw_only=True) -class Cosmos2I2VRunnerConfig(Cosmos2T2VRunnerConfig): - """Runner config for the Cosmos-Predict2 I2V variant.""" - - _target: type["Cosmos2I2VRunner"] = field(default_factory=lambda: Cosmos2I2VRunner) - - image_path: str | Path = DEFAULT_I2V_IMAGE_URL - """First-frame RGB image. Either a local path or an HTTP(S) URL.""" - - -class Cosmos2I2VRunner(Cosmos2T2VRunner): - """Cosmos-Predict2 non-streaming I2V driver.""" - - config: Cosmos2I2VRunnerConfig - - def _initialize_cache(self) -> CosmosInferencePipelineCache: - """Initialize the autoregressive cache for I2V (loads first frame).""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - - # Load + resize the first frame, then convert to [-1, 1] bf16 - # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the - # pipeline's actual device so non-default ``--device`` selections - # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - image = load_first_frame_tensor( - resolve_input_path( - config.image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=read_image_rgb, - ), - pixel_height=config.pixel_height, - pixel_width=config.pixel_width, - device=self.pipeline.device, - dtype=torch.bfloat16, - ) - - return self.pipeline.initialize_cache(text=[prompt], image=image) diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/config.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/config.py index c7b819c85..66e3753d4 100644 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/config.py +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/config.py @@ -17,9 +17,38 @@ from __future__ import annotations +from dataclasses import dataclass +from pathlib import Path + import torch -from fastvideo_causal_wan22.runner import FastvideoCausalWan22T2VRunnerConfig +from flashdreams.runtime.video_runner import StreamingVideoRunnerConfig + +DEFAULT_T2V_PROMPT = ( + "A stylish woman strolls down a bustling Tokyo street, the warm glow of " + "neon lights and animated city signs casting vibrant reflections. She " + "wears a sleek black leather jacket paired with a flowing red dress and " + "black boots, her black purse slung over her shoulder. Sunglasses perched " + "on her nose and a bold red lipstick add to her confident, casual " + "demeanor. The street is damp and reflective, creating a mirror-like " + "effect that enhances the colorful lights and shadows. Pedestrians move " + "about, adding to the lively atmosphere. The scene is captured in a " + "dynamic medium shot with the woman walking slightly to one side, " + "highlighting her graceful strides." +) + + +@dataclass(kw_only=True) +class FastvideoCausalWan22T2VRunnerConfig(StreamingVideoRunnerConfig): + """Runner config for the FastVideo CausalWan 2.2 T2V variants.""" + + prompt: str | Path = DEFAULT_T2V_PROMPT + total_blocks: int = 60 + pixel_height: int = 480 + pixel_width: int = 832 + fps: int = 16 + + from flashdreams.core.checkpoint.remap import remap_checkpoint_keys from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py deleted file mode 100644 index 781ca58e5..000000000 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py +++ /dev/null @@ -1,174 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""FastVideo CausalWan 2.2 streaming T2V runner class.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path - -from loguru import logger - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) -from flashdreams.recipes.wan import ( - WanInferencePipeline, - WanInferencePipelineCache, -) -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - -__all__ = [ - "FastvideoCausalWan22T2VRunnerConfig", - "FastvideoCausalWan22T2VRunner", -] - - -DEFAULT_T2V_PROMPT = ( - "A stylish woman strolls down a bustling Tokyo street, the warm glow of " - "neon lights and animated city signs casting vibrant reflections. She " - "wears a sleek black leather jacket paired with a flowing red dress and " - "black boots, her black purse slung over her shoulder. Sunglasses perched " - "on her nose and a bold red lipstick add to her confident, casual " - "demeanor. The street is damp and reflective, creating a mirror-like " - "effect that enhances the colorful lights and shadows. Pedestrians move " - "about, adding to the lively atmosphere. The scene is captured in a " - "dynamic medium shot with the woman walking slightly to one side, " - "highlighting her graceful strides." -) - - -@dataclass(kw_only=True) -class FastvideoCausalWan22T2VRunnerConfig(RunnerConfig): - """Runner config for the FastVideo CausalWan 2.2 T2V variants.""" - - _target: type["FastvideoCausalWan22T2VRunner"] = field( - default_factory=lambda: FastvideoCausalWan22T2VRunner - ) - - prompt: str | Path = DEFAULT_T2V_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt).""" - - total_blocks: int = 60 - """Number of autoregressive chunks to generate before terminating the rollout.""" - - pixel_height: int = 480 - """Output video pixel height.""" - - pixel_width: int = 832 - """Output video pixel width.""" - - fps: int = 16 - """Output video frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Pipeline output layout for streaming post-processing.""" - - -class FastvideoCausalWan22T2VRunner( - Runner[FastvideoCausalWan22T2VRunnerConfig, WanInferencePipeline] -): - """FastVideo CausalWan 2.2 streaming T2V driver (14B MoE, 8-step distilled).""" - - config: FastvideoCausalWan22T2VRunnerConfig - - def _resolve_prompt(self) -> str: - """Resolve config.prompt. - - A Path reads its first non-empty line, a str is used as-is. - """ - return resolve_prompt_value(self.config.prompt) - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache.""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - spatial_compression_ratio = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % spatial_compression_ratio == 0, ( - f"pixel_height={self.config.pixel_height} must divide " - f"{spatial_compression_ratio}." - ) - assert config.pixel_width % spatial_compression_ratio == 0, ( - f"pixel_width={self.config.pixel_width} must divide {spatial_compression_ratio}." - ) - latent_h = config.pixel_height // spatial_compression_ratio - latent_w = config.pixel_width // spatial_compression_ratio - - return self.pipeline.initialize_cache( - text=[prompt], image=None, height=latent_h, width=latent_w - ) - - def run(self) -> None: - """Drive the autoregressive rollout and write outputs.""" - config = self.config - - # Initialize the autoregressive cache. - cache = self._initialize_cache() - - output_stream = self.create_video_output_stream(fps=config.fps) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - for i in range(config.total_blocks): - # Generate the autoregressive chunks. - video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_target.write( - output_stream.process( - video_chunk, - autoregressive_index=i, - metrics=stats, - ) - ) - - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - video_artifact = artifacts[0] - video_path = Path(video_artifact.uri) - - logger.info( - f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " - f"-> {video_path.resolve()}" - ) - - # Write the perf stats. - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) diff --git a/integrations/self_forcing/self_forcing/config.py b/integrations/self_forcing/self_forcing/config.py index becf1ecd5..9c69e127c 100644 --- a/integrations/self_forcing/self_forcing/config.py +++ b/integrations/self_forcing/self_forcing/config.py @@ -17,6 +17,8 @@ from __future__ import annotations +from dataclasses import dataclass +from pathlib import Path from typing import Any, cast from torch import Tensor @@ -32,7 +34,32 @@ WanInferencePipelineConfig, WanVAEDecoderConfig, ) -from self_forcing.runner import SelfForcingT2VRunnerConfig +from flashdreams.runtime.video_runner import StreamingVideoRunnerConfig + +DEFAULT_T2V_PROMPT = ( + "A stylish woman strolls down a bustling Tokyo street, the warm glow of " + "neon lights and animated city signs casting vibrant reflections. She " + "wears a sleek black leather jacket paired with a flowing red dress and " + "black boots, her black purse slung over her shoulder. Sunglasses perched " + "on her nose and a bold red lipstick add to her confident, casual " + "demeanor. The street is damp and reflective, creating a mirror-like " + "effect that enhances the colorful lights and shadows. Pedestrians move " + "about, adding to the lively atmosphere. The scene is captured in a " + "dynamic medium shot with the woman walking slightly to one side, " + "highlighting her graceful strides." +) + + +@dataclass(kw_only=True) +class SelfForcingT2VRunnerConfig(StreamingVideoRunnerConfig): + """Runner config for the Self-Forcing T2V variants.""" + + prompt: str | Path = DEFAULT_T2V_PROMPT + total_blocks: int = 60 + pixel_height: int = 480 + pixel_width: int = 832 + fps: int = 16 + CHECKPOINT_PATH = "https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt" diff --git a/integrations/self_forcing/self_forcing/runner.py b/integrations/self_forcing/self_forcing/runner.py deleted file mode 100644 index 3e4857c7f..000000000 --- a/integrations/self_forcing/self_forcing/runner.py +++ /dev/null @@ -1,172 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Self-Forcing Wan 2.1 streaming T2V runner class.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path - -from loguru import logger - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) -from flashdreams.recipes.wan import ( - WanInferencePipeline, - WanInferencePipelineCache, -) -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - -__all__ = [ - "SelfForcingT2VRunnerConfig", - "SelfForcingT2VRunner", -] - - -DEFAULT_T2V_PROMPT = ( - "A stylish woman strolls down a bustling Tokyo street, the warm glow of " - "neon lights and animated city signs casting vibrant reflections. She " - "wears a sleek black leather jacket paired with a flowing red dress and " - "black boots, her black purse slung over her shoulder. Sunglasses perched " - "on her nose and a bold red lipstick add to her confident, casual " - "demeanor. The street is damp and reflective, creating a mirror-like " - "effect that enhances the colorful lights and shadows. Pedestrians move " - "about, adding to the lively atmosphere. The scene is captured in a " - "dynamic medium shot with the woman walking slightly to one side, " - "highlighting her graceful strides." -) - - -@dataclass(kw_only=True) -class SelfForcingT2VRunnerConfig(RunnerConfig): - """Runner config for the Self-Forcing T2V variants.""" - - _target: type["SelfForcingT2VRunner"] = field( - default_factory=lambda: SelfForcingT2VRunner - ) - - prompt: str | Path = DEFAULT_T2V_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt).""" - - total_blocks: int = 60 - """Number of autoregressive chunks to generate before terminating the rollout.""" - - pixel_height: int = 480 - """Output video pixel height.""" - - pixel_width: int = 832 - """Output video pixel width.""" - - fps: int = 16 - """Output video frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Pipeline output layout for streaming post-processing.""" - - -class SelfForcingT2VRunner(Runner[SelfForcingT2VRunnerConfig, WanInferencePipeline]): - """Self-Forcing Wan 2.1 streaming T2V driver.""" - - config: SelfForcingT2VRunnerConfig - - def _resolve_prompt(self) -> str: - """Resolve config.prompt. - - A Path reads its first non-empty line, a str is used as-is. - """ - return resolve_prompt_value(self.config.prompt) - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache.""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - spatial_compression_ratio = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % spatial_compression_ratio == 0, ( - f"pixel_height={self.config.pixel_height} must divide " - f"{spatial_compression_ratio}." - ) - assert config.pixel_width % spatial_compression_ratio == 0, ( - f"pixel_width={self.config.pixel_width} must divide {spatial_compression_ratio}." - ) - latent_h = config.pixel_height // spatial_compression_ratio - latent_w = config.pixel_width // spatial_compression_ratio - - return self.pipeline.initialize_cache( - text=[prompt], image=None, height=latent_h, width=latent_w - ) - - def run(self) -> None: - """Drive the autoregressive rollout and write outputs.""" - config = self.config - - # Initialize the autoregressive cache. - cache = self._initialize_cache() - - # Generate the autoregressive chunks. - output_stream = self.create_video_output_stream(fps=config.fps) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - for i in range(config.total_blocks): - video_chunk = self.pipeline.generate(autoregressive_index=i, cache=cache) - stats = self.pipeline.finalize(autoregressive_index=i, cache=cache) - output_target.write( - output_stream.process( - video_chunk, - autoregressive_index=i, - metrics=stats, - ) - ) - - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - video_artifact = artifacts[0] - video_path = Path(video_artifact.uri) - - logger.info( - f"[{config.runner_name}] wrote video {video_artifact.metadata['shape']} " - f"-> {video_path.resolve()}" - ) - - # Write the perf stats. - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) diff --git a/integrations/wan21/wan21/config.py b/integrations/wan21/wan21/config.py index 717cbfa26..a1cdacdc7 100644 --- a/integrations/wan21/wan21/config.py +++ b/integrations/wan21/wan21/config.py @@ -17,6 +17,9 @@ from __future__ import annotations +from dataclasses import dataclass +from pathlib import Path + from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler import ( FlowMatchUniPCSchedulerConfig, @@ -32,7 +35,46 @@ WanVAEDecoderConfig, WanVAEEncoderConfig, ) -from wan21.runner import Wan21I2VRunnerConfig, Wan21T2VRunnerConfig +from flashdreams.runtime.video_runner import ( + ImageConditionedVideoRunnerConfig, + VideoRunnerConfig, +) + +DEFAULT_PROMPT = ( + "Summer beach vacation style, a white cat wearing sunglasses sits on " + "a surfboard. The fluffy-furred feline gazes directly at the camera " + "with a relaxed expression. Blurred beach scenery forms the background " + "featuring crystal-clear waters, distant green hills, and a blue sky " + "dotted with white clouds. The cat assumes a naturally relaxed posture, " + "as if savoring the sea breeze and warm sunlight. A close-up shot " + "highlights the feline's intricate details and the refreshing " + "atmosphere of the seaside." +) + +DEFAULT_I2V_IMAGE_URL = ( + "https://raw.githubusercontent.com/Wan-Video/Wan2.1/main/examples/i2v_input.JPG" +) + + +@dataclass(kw_only=True) +class Wan21T2VRunnerConfig(VideoRunnerConfig): + """Runner config for the Wan 2.1 T2V variant.""" + + prompt: str | Path = DEFAULT_PROMPT + pixel_height: int = 480 + pixel_width: int = 832 + fps: int = 16 + + +@dataclass(kw_only=True) +class Wan21I2VRunnerConfig(ImageConditionedVideoRunnerConfig, Wan21T2VRunnerConfig): + """Runner config for the Wan 2.1 I2V variant.""" + + image_path: str | Path = DEFAULT_I2V_IMAGE_URL + image_cache_subdir = "wan21" + pixel_height: int = 832 + pixel_width: int = 480 + CHECKPOINT_PATH_T2V_1PT3B = ( "https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B/blob/main/" diff --git a/integrations/wan21/wan21/runner.py b/integrations/wan21/wan21/runner.py deleted file mode 100644 index 2098c746b..000000000 --- a/integrations/wan21/wan21/runner.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Non-streaming Wan 2.1 runner classes (T2V and I2V).""" - -from __future__ import annotations - -import os -from dataclasses import dataclass, field -from pathlib import Path - -import torch -from loguru import logger - -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.infra.runner import Runner, RunnerConfig -from flashdreams.infra.runner_io import ( - load_first_frame_tensor, - read_image_rgb, - resolve_input_path, - resolve_prompt_value, - runner_artifact_path, - write_runner_stats, -) -from flashdreams.recipes.wan import ( - WanInferencePipeline, - WanInferencePipelineCache, -) -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - -__all__ = [ - "Wan21I2VRunnerConfig", - "Wan21I2VRunner", - "Wan21T2VRunnerConfig", - "Wan21T2VRunner", -] - - -DEFAULT_PROMPT = ( - "Summer beach vacation style, a white cat wearing sunglasses sits on " - "a surfboard. The fluffy-furred feline gazes directly at the camera " - "with a relaxed expression. Blurred beach scenery forms the background " - "featuring crystal-clear waters, distant green hills, and a blue sky " - "dotted with white clouds. The cat assumes a naturally relaxed posture, " - "as if savoring the sea breeze and warm sunlight. A close-up shot " - "highlights the feline's intricate details and the refreshing " - "atmosphere of the seaside." -) - -DEFAULT_I2V_IMAGE_URL = ( - "https://raw.githubusercontent.com/Wan-Video/Wan2.1/main/examples/i2v_input.JPG" -) - -IMAGE_CACHE_DIR = ( - Path(os.path.expanduser(os.getenv("FLASHDREAMS_CACHE_DIR", "~/.cache/flashdreams"))) - / "wan21" -) -"""User-writable cache for on-the-fly I2V first-frame downloads.""" - - -@dataclass(kw_only=True) -class Wan21T2VRunnerConfig(RunnerConfig): - """Runner config for the Wan 2.1 T2V variant. - - Also serves as the base for :class:`Wan21I2VRunnerConfig` - (I2V is T2V plus an ``image_path``). - """ - - _target: type["Wan21T2VRunner"] = field(default_factory=lambda: Wan21T2VRunner) - - prompt: str | Path = DEFAULT_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt). - Defaults to :data:`DEFAULT_PROMPT`.""" - - pixel_height: int = 480 - """Output video pixel height.""" - - pixel_width: int = 832 - """Output video pixel width.""" - - fps: int = 16 - """Output video frame rate.""" - - postprocess_output_layout: VideoTensorLayout | None = "tchw" - """Pipeline output layout for streaming post-processing.""" - - -@dataclass(kw_only=True) -class Wan21I2VRunnerConfig(Wan21T2VRunnerConfig): - """Runner config for the Wan 2.1 I2V variant. - - Inherits all T2V fields (prompt, pixel_*, fps) and - adds the first-frame image path that I2V needs at runtime. - """ - - _target: type["Wan21I2VRunner"] = field(default_factory=lambda: Wan21I2VRunner) - - image_path: str | Path = DEFAULT_I2V_IMAGE_URL - """Path to the first-frame RGB image, or an ``http(s)://`` URL that - will be downloaded on first use into :data:`IMAGE_CACHE_DIR`. - Defaults to :data:`DEFAULT_I2V_IMAGE_URL`.""" - - prompt: str | Path = DEFAULT_PROMPT - """Either an inline text prompt (--prompt "...") or a path to a - txt file whose first line is read as the prompt (--prompt prompt.txt). - Defaults to :data:`DEFAULT_PROMPT`.""" - - pixel_height: int = 832 - """Output video pixel height.""" - - pixel_width: int = 480 - """Output video pixel width.""" - - -class Wan21T2VRunner(Runner[Wan21T2VRunnerConfig, WanInferencePipeline]): - """Wan 2.1 non-streaming T2V driver. - - Also serves as the base for :class:`Wan21I2VRunner` (I2V - only overrides :meth:`_initialize_cache` to load the first frame; - everything else, including :meth:`run`, is reused). - """ - - config: Wan21T2VRunnerConfig - - def _resolve_prompt(self) -> str: - """Resolve config.prompt. - - A Path reads its first non-empty line, a str is used as-is. - """ - return resolve_prompt_value(self.config.prompt) - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache for T2V.""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - latent_h = config.pixel_height // sp - latent_w = config.pixel_width // sp - - return self.pipeline.initialize_cache( - text=[prompt], image=None, height=latent_h, width=latent_w - ) - - def run(self) -> None: - """Drive the single-step rollout and write outputs.""" - config = self.config - - # Initialize the autoregressive cache. - cache = self._initialize_cache() - - # Generate the output in one AR step. - output_stream = self.create_video_output_stream(fps=config.fps) - video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") - output_target = Mp4VideoOutputTarget( - output_path=video_path, - fps=config.fps, - output_layout=output_stream.output_layout, - enabled=self.is_rank_zero, - ) - output_target.open() - generated = self.pipeline.generate(autoregressive_index=0, cache=cache) - stats = self.pipeline.finalize(autoregressive_index=0, cache=cache) - output_target.write( - output_stream.process(generated, autoregressive_index=0, metrics=stats) - ) - tail = output_stream.finish() - if tail is not None: - output_target.write(tail) - artifacts = output_target.close() - if not artifacts: - return - video_artifact = artifacts[0] - video_path = Path(video_artifact.uri) - - logger.info( - f"[{config.runner_name}] wrote video {tuple(generated.shape)} " - f"-> {video_path.resolve()}" - ) - - # Write the perf stats. - stats_history = video_artifact.metadata["stats_history"] - if stats_history: - stats_path = write_runner_stats( - config.output_dir, - config.runner_name, - list(stats_history), - ) - logger.info( - f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" - ) - - -class Wan21I2VRunner(Wan21T2VRunner): - """Wan 2.1 non-streaming I2V driver (first-frame injection).""" - - config: Wan21I2VRunnerConfig - - def _initialize_cache(self) -> WanInferencePipelineCache: - """Initialize the autoregressive cache for I2V (loads first frame).""" - config = self.config - prompt = self._resolve_prompt() - - assert isinstance(self.pipeline.decoder, StreamingVideoDecoder) - sp = self.pipeline.decoder.spatial_compression_ratio - assert config.pixel_height % sp == 0, ( - f"pixel_height={config.pixel_height} must divide {sp}." - ) - assert config.pixel_width % sp == 0, ( - f"pixel_width={config.pixel_width} must divide {sp}." - ) - - # Load + resize the first frame, then convert to [-1, 1] bf16 - # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the - # pipeline's actual device so non-default ``--device`` selections - # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - image = load_first_frame_tensor( - resolve_input_path( - config.image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=read_image_rgb, - ), - pixel_height=config.pixel_height, - pixel_width=config.pixel_width, - device=self.pipeline.device, - dtype=torch.bfloat16, - ) - - return self.pipeline.initialize_cache(text=[prompt], image=image) diff --git a/uv.lock b/uv.lock index c37b144be..5065e7bc7 100644 --- a/uv.lock +++ b/uv.lock @@ -1430,11 +1430,11 @@ name = "flashdreams-t2v" version = "0.1.0" source = { editable = "apps/t2v" } dependencies = [ - { name = "flashdreams" }, + { name = "flashdreams", extra = ["serving"] }, ] [package.metadata] -requires-dist = [{ name = "flashdreams", editable = "flashdreams" }] +requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] [[package]] name = "flashdreams-t2v-demo" From 223bf4b54da30edb1a0372ab7a823d367cb5dffc Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 17:23:33 +0000 Subject: [PATCH 27/33] Clean up direct WebRTC app startup failures Ensure direct public WebRTC app launches close the created runtime when startup fails before aiohttp owns shutdown. Route failure cleanup through the WebRTC manager and RuntimeHost path, preserving the original startup error, and add CPU coverage for the cleanup behavior. --- apps/t2v/tests/test_t2v_shell.py | 1 - apps/t2v_demo/README.md | 22 - apps/t2v_demo/__init__.py | 4 - apps/t2v_demo/app.py | 382 ------------------ apps/t2v_demo/launch.py | 88 ---- apps/t2v_demo/pyproject.toml | 39 -- apps/t2v_demo/runner.py | 70 ---- apps/t2v_demo/runtime.py | 226 ----------- apps/t2v_demo/tests/test_runner.py | 133 ------ apps/t2v_demo/web/adapter.css | 3 - apps/t2v_demo/web/adapter.js | 23 -- flashdreams/flashdreams/demo/app.py | 73 +++- .../tests/test_demo_application_api.py | 37 ++ uv.lock | 22 - 14 files changed, 89 insertions(+), 1034 deletions(-) delete mode 100644 apps/t2v_demo/README.md delete mode 100644 apps/t2v_demo/__init__.py delete mode 100644 apps/t2v_demo/app.py delete mode 100644 apps/t2v_demo/launch.py delete mode 100644 apps/t2v_demo/pyproject.toml delete mode 100644 apps/t2v_demo/runner.py delete mode 100644 apps/t2v_demo/runtime.py delete mode 100644 apps/t2v_demo/tests/test_runner.py delete mode 100644 apps/t2v_demo/web/adapter.css delete mode 100644 apps/t2v_demo/web/adapter.js diff --git a/apps/t2v/tests/test_t2v_shell.py b/apps/t2v/tests/test_t2v_shell.py index b03e6577a..666968fbd 100644 --- a/apps/t2v/tests/test_t2v_shell.py +++ b/apps/t2v/tests/test_t2v_shell.py @@ -136,7 +136,6 @@ def test_t2v_shell_has_no_legacy_backend_imports() -> None: source = inspect.getsource(t2v_shell) - assert "t2v_demo" not in source assert "causal_forcing" not in source assert "self_forcing" not in source assert "cosmos_predict2" not in source diff --git a/apps/t2v_demo/README.md b/apps/t2v_demo/README.md deleted file mode 100644 index c29d847a6..000000000 --- a/apps/t2v_demo/README.md +++ /dev/null @@ -1,22 +0,0 @@ -# FlashDreams text-to-video demo - -This app uses the shared `flashdreams.runtime.demo` replay/session lifecycle. -The selected integration owns its pipeline and checkpoint configuration; the -app owns only prompt input, output selection, and the thin runtime adapter. - -Run a saved replay: - -```bash -uv run python -m apps.t2v_demo.app replay --backend causal-forcing --output outputs/t2v.mp4 -``` - -Serve streamed WebRTC output: - -```bash -uv run python -m apps.t2v_demo.app webrtc --backend self-forcing -``` - -`--backend` accepts `causal-forcing`, `cosmos-predict2`, and `self-forcing`. -Use `--preset-id` to select an integration-specific T2V runner preset. The -browser UI accepts a prompt before opening a generation session, plays emitted -chunks as they arrive, and records the received WebRTC stream for download. diff --git a/apps/t2v_demo/__init__.py b/apps/t2v_demo/__init__.py deleted file mode 100644 index 741c93237..000000000 --- a/apps/t2v_demo/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Unified text-to-video demo backed by the FlashDreams runtime API.""" diff --git a/apps/t2v_demo/app.py b/apps/t2v_demo/app.py deleted file mode 100644 index 498374d99..000000000 --- a/apps/t2v_demo/app.py +++ /dev/null @@ -1,382 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Typed ``flashdreams-run t2v`` launch implementation.""" - -from __future__ import annotations - -import io -import json -import zipfile -from dataclasses import dataclass, replace -from importlib.resources import files -from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal - -from aiohttp import web -from t2v.t2v import ( - T2VRunDefaults, - create_t2v_application, - create_t2v_spec, - t2v_scenario_mapping, -) - -from flashdreams.demo import Application, run_application_replay -from flashdreams.runtime.demo import ( - DemoSpec, - Mp4OutputSpec, - NullOutputSpec, - WebRTCAppResources, - WebRTCOutputSpec, -) -from flashdreams.runtime.demo.bootstrap import ( - configure_logging, - initialize_cuda_distributed, -) -from flashdreams.runtime.demo.host import RuntimeHost -from flashdreams.serving.webrtc.demo import serve_webrtc_demo -from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager -from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig - -from .runtime import ( - FIELD_FPS, - FIELD_PIXEL_HEIGHT, - FIELD_PIXEL_WIDTH, - FIELD_PROMPT, - FIELD_TOTAL_BLOCKS, - T2VDemoAdapter, - backend_metadata, - make_adapter, - model_from_backend, -) - -if TYPE_CHECKING: - from .runner import T2VDemoRunnerConfig - - -@dataclass(frozen=True, slots=True) -class T2VWebRTCConfig(WebRTCRuntimeConfig): - """Shared WebRTC settings required by the prompt-only T2V demo.""" - - video_width: int - video_height: int - warmup_chunks: int - warmup_timeout_s: float - - -@dataclass(frozen=True, kw_only=True, slots=True) -class T2VApplicationDefaults(T2VRunDefaults): - """Legacy ``t2v`` command defaults plus temporary backend selection.""" - - backend: str = "causal-forcing" - preset_id: str | None = None - - -class T2VWebRTCSessionManager(BaseWebRTCSessionManager[Any, T2VWebRTCConfig]): - """Shared manager with a prompt update for the next browser session.""" - - def update_prompt(self, prompt: str, duration_s: float) -> None: - if not prompt.strip(): - raise ValueError("Prompt must be non-empty.") - if not 0 < duration_s <= 60: - raise ValueError("Duration must be greater than 0 and at most 60 seconds.") - spec = self._shared_spec - adapter = self._shared_adapter - if spec is None or adapter is None: - raise RuntimeError("T2V WebRTC shared session is not initialized.") - scenario = dict(spec.scenario or {}) - scenario[FIELD_PROMPT] = prompt.strip() - scenario[FIELD_TOTAL_BLOCKS] = self.runtime.blocks_for_duration( - duration_s, fps=_int_value(scenario[FIELD_FPS], name=FIELD_FPS) - ) - spec = replace(spec, scenario=scenario) - self._shared_spec = spec - self._shared_scenario = adapter.prepare_scenario(spec) - - -def launch_t2v( - *, - config: "T2VDemoRunnerConfig", - mode: Literal["mp4", "null", "webrtc"], - scenario_overrides: dict[str, object] | None = None, - output_overrides: dict[str, object] | None = None, - host: str | None = None, - port: int | None = None, -) -> object: - """Launch T2V directly from its typed ``flashdreams-run`` configuration.""" - configure_logging() - scenario_overrides = scenario_overrides or {} - output_overrides = output_overrides or {} - defaults = _defaults_from_config(config, scenario_overrides) - scenario = _scenario(defaults) - if mode == "mp4" or mode == "null": - output = _replay_output( - mode=mode, - output_path=output_overrides.get( - "path", output_overrides.get("output", config.output) - ), - fps=_int_value( - output_overrides.get("fps", scenario[FIELD_FPS]), name="fps" - ), - ) - result = run_application_replay( - app=create_t2v_application( - model=model_from_backend(defaults.backend, defaults.preset_id), - defaults=defaults, - output=output, - ) - ) - if result.status != "completed": - reason = result.reason or str(result.error) or "T2V replay failed." - raise RuntimeError(reason) - return result - - context = initialize_cuda_distributed(default_device=config.device) - # The browser download and playback endpoints below read the session - # artifact this enables. - adapter = make_adapter( - defaults.backend, defaults.preset_id, write_download_artifact=True - ) - output = WebRTCOutputSpec( - host=str(host or output_overrides.get("host", "0.0.0.0")), - port=_int_value( - port if port is not None else output_overrides.get("port", 8080), - name="port", - ), - fps=_int_value(output_overrides.get("fps", scenario[FIELD_FPS]), name="fps"), - video_width=_int_value( - output_overrides.get("video_width", scenario[FIELD_PIXEL_WIDTH]), - name="video_width", - ), - video_height=_int_value( - output_overrides.get("video_height", scenario[FIELD_PIXEL_HEIGHT]), - name="video_height", - ), - warmup_chunks=_int_value( - output_overrides.get("warmup_chunks", 0), name="warmup_chunks" - ), - warmup_timeout_s=_float_value( - output_overrides.get("warmup_timeout_s", 600.0), - name="warmup_timeout_s", - ), - client_liveness_timeout_s=_float_value( - output_overrides.get("client_liveness_timeout_s", 30.0), - name="client_liveness_timeout_s", - ), - preload_name="FlashDreams T2V", - ) - spec = _spec( - replace(defaults, device=str(context.device)), - adapter=adapter, - input_mode="webrtc", - output=output, - ) - prepared = adapter.prepare_scenario(spec) - inference_config = spec.config - if inference_config is None: - raise RuntimeError("T2V DemoSpec.config was not initialized.") - runtime = adapter.create_runtime(inference_config) - manager = T2VWebRTCSessionManager( - runtime=runtime, - runtime_config=T2VWebRTCConfig( - video_width=output.video_width, - video_height=output.video_height, - warmup_chunks=output.warmup_chunks, - warmup_timeout_s=output.warmup_timeout_s, - ), - fps=output.fps, - identity=adapter.model_id, - supported_control_keys=frozenset({"g"}), - shared_host=RuntimeHost(runtime), - shared_adapter=adapter, - shared_spec=spec, - shared_scenario=prepared, - client_liveness_timeout_s=output.client_liveness_timeout_s, - keep_connection_after_completed=True, - ) - return serve_webrtc_demo( - output=output, - model_id=adapter.model_id, - session_manager=manager, - app_resources=WebRTCAppResources( - model_web_resource=files("t2v_demo").joinpath("web"), - configure_app=lambda app: _configure_app( - app, manager=manager, backend=config.backend - ), - preload_name="FlashDreams T2V", - ), - world_rank=context.world_rank, - ) - - -def create_app(config: "T2VDemoRunnerConfig | None" = None) -> Application: - """Create the default public T2V application without CLI parsing.""" - from .runner import RUNNER_T2V - - config = RUNNER_T2V if config is None else config - defaults = _defaults_from_config(config, {}) - return create_t2v_application( - model=model_from_backend(defaults.backend, defaults.preset_id), - defaults=defaults, - ) - - -createApp = create_app - - -def _defaults_from_config( - config: "T2VDemoRunnerConfig", - overrides: dict[str, object], -) -> T2VApplicationDefaults: - def value(name: str) -> object: - return getattr(config, name) if overrides.get(name) is None else overrides[name] - - return T2VApplicationDefaults( - backend=config.backend, - preset_id=config.preset_id, - prompt=_optional_str(value(FIELD_PROMPT)), - total_blocks=_optional_int(value(FIELD_TOTAL_BLOCKS)), - pixel_height=_optional_int(value(FIELD_PIXEL_HEIGHT)), - pixel_width=_optional_int(value(FIELD_PIXEL_WIDTH)), - fps=_optional_int(value(FIELD_FPS)), - device=config.device, - compile=config.compile, - ) - - -def _scenario(defaults: T2VApplicationDefaults) -> dict[str, object]: - return t2v_scenario_mapping( - model=model_from_backend(defaults.backend, defaults.preset_id), - defaults=defaults, - ) - - -def _spec( - defaults: T2VApplicationDefaults, - *, - adapter: T2VDemoAdapter, - input_mode: Literal["replay", "webrtc"], - output: Mp4OutputSpec | NullOutputSpec | WebRTCOutputSpec, -) -> DemoSpec: - return create_t2v_spec( - model=adapter.model, - defaults=defaults, - input_mode=input_mode, - output=output, - ) - - -def _replay_output( - *, mode: Literal["mp4", "null"], output_path: object, fps: int -) -> Mp4OutputSpec | NullOutputSpec: - if mode == "null": - return NullOutputSpec() - if output_path is None: - raise ValueError("T2V MP4 mode requires an output path.") - return Mp4OutputSpec(path=Path(str(output_path)), fps=fps, output_layout="tchw") - - -def _int_value(value: object, *, name: str) -> int: - if isinstance(value, bool): - raise TypeError(f"{name} must be an integer, not bool.") - if isinstance(value, int): - return value - if isinstance(value, float): - return int(value) - if isinstance(value, str): - return int(value) - raise TypeError(f"{name} must be convertible to int, got {type(value).__name__}.") - - -def _float_value(value: object, *, name: str) -> float: - if isinstance(value, bool): - raise TypeError(f"{name} must be numeric, not bool.") - if isinstance(value, int | float): - return float(value) - if isinstance(value, str): - return float(value) - raise TypeError(f"{name} must be convertible to float, got {type(value).__name__}.") - - -def _optional_int(value: object) -> int | None: - return None if value is None else _int_value(value, name="value") - - -def _optional_str(value: object) -> str | None: - return None if value is None else str(value) - - -def _configure_app( - app: web.Application, - *, - manager: T2VWebRTCSessionManager, - backend: str, -) -> None: - async def app_config(_: web.Request) -> web.StreamResponse: - return web.json_response( - {"backends": backend_metadata(), "selected_backend": backend} - ) - - async def update_prompt(request: web.Request) -> web.StreamResponse: - payload = await request.json() - if not isinstance(payload, dict) or not isinstance(payload.get("prompt"), str): - raise web.HTTPBadRequest(reason="Expected a JSON prompt.") - duration_s = payload.get("duration_s") - if not isinstance(duration_s, int | float): - raise web.HTTPBadRequest(reason="Expected numeric duration_s.") - try: - manager.update_prompt(payload["prompt"], float(duration_s)) - except (RuntimeError, ValueError) as exc: - raise web.HTTPBadRequest(reason=str(exc)) from exc - return web.json_response({"status": "ok"}) - - async def download(_: web.Request) -> web.StreamResponse: - artifact = manager.runtime.latest_artifact - if artifact is None: - raise web.HTTPNotFound(reason="No completed generation is available yet.") - video_path, scenario = artifact - if not video_path.is_file(): - raise web.HTTPNotFound(reason="Generated MP4 is no longer available.") - buffer = io.BytesIO() - with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: - archive.write(video_path, "video.mp4") - archive.writestr( - "prompt.json", - json.dumps( - { - "prompt": scenario.prompt, - "total_blocks": scenario.total_blocks, - "fps": scenario.fps, - "width": scenario.pixel_width, - "height": scenario.pixel_height, - }, - indent=2, - ), - ) - return web.Response( - body=buffer.getvalue(), - headers={ - "Content-Disposition": "attachment; filename=flashdreams-generation.zip" - }, - content_type="application/zip", - ) - - async def playback(_: web.Request) -> web.StreamResponse: - artifact = manager.runtime.latest_artifact - if artifact is None or not artifact[0].is_file(): - raise web.HTTPNotFound(reason="No completed MP4 is available yet.") - return web.FileResponse(artifact[0]) - - app.router.add_get("/api/t2v/config", app_config) - app.router.add_post("/api/t2v/prompt", update_prompt) - app.router.add_get("/api/t2v/download", download) - app.router.add_get("/api/t2v/playback", playback) - - -__all__ = [ - "T2VApplicationDefaults", - "T2VWebRTCSessionManager", - "createApp", - "create_app", - "launch_t2v", -] diff --git a/apps/t2v_demo/launch.py b/apps/t2v_demo/launch.py deleted file mode 100644 index 3b818a6c7..000000000 --- a/apps/t2v_demo/launch.py +++ /dev/null @@ -1,88 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Launch capability that routes ``flashdreams-run t2v`` to the T2V app.""" - -from __future__ import annotations - -from functools import partial -from typing import Literal, TypeAlias - -from flashdreams.infra.runner import RunnerConfig -from flashdreams.serving.launch import LaunchMode, LaunchOptions, ResolvedLaunch - -from .runner import T2VDemoRunnerConfig - -T2VLaunchMode: TypeAlias = Literal["mp4", "null", "webrtc"] - - -class T2VLaunchCapability: - """Expose replay and persistent WebRTC modes for the app-owned T2V demo.""" - - def supported_modes( - self, config: RunnerConfig, options: LaunchOptions - ) -> tuple[LaunchMode, ...]: - del config, options - return ("mp4", "null", "webrtc") - - def resolve( - self, - config: RunnerConfig, - *, - mode: LaunchMode, - options: LaunchOptions, - ) -> ResolvedLaunch | None: - t2v_mode = _t2v_mode(mode) - if t2v_mode is None: - return None - t2v_config = _t2v_config(config) - return ResolvedLaunch( - mode=t2v_mode, - label=f"T2V {t2v_mode} launch", - summary={ - "runner": t2v_config.runner_name, - "mode": t2v_mode, - "device": t2v_config.device, - }, - launch=partial( - _launch, - config=t2v_config, - mode=t2v_mode, - options=options, - ), - ) - - -def _launch( - *, config: T2VDemoRunnerConfig, mode: T2VLaunchMode, options: LaunchOptions -) -> object: - from .app import launch_t2v - - return launch_t2v( - config=config, - mode=mode, - host=options.host, - port=options.port, - scenario_overrides=dict(options.scenario), - output_overrides=dict(options.output), - ) - - -def _t2v_mode(mode: LaunchMode) -> T2VLaunchMode | None: - if mode == "mp4" or mode == "null" or mode == "webrtc": - return mode - return None - - -def _t2v_config(config: RunnerConfig) -> T2VDemoRunnerConfig: - if not isinstance(config, T2VDemoRunnerConfig): - raise TypeError( - "T2V launch capability requires T2VDemoRunnerConfig, got " - f"{type(config).__name__}." - ) - return config - - -LAUNCH_CAPABILITY = T2VLaunchCapability() - -__all__ = ["LAUNCH_CAPABILITY", "T2VLaunchCapability"] diff --git a/apps/t2v_demo/pyproject.toml b/apps/t2v_demo/pyproject.toml deleted file mode 100644 index 64453965c..000000000 --- a/apps/t2v_demo/pyproject.toml +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -[build-system] -requires = ["setuptools>=69", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "flashdreams-t2v-demo" -version = "0.1.0" -description = "FlashDreams text-to-video runtime demo launcher" -requires-python = ">=3.10" -dependencies = [ - "flashdreams[serving]", - "flashdreams-causal-forcing", - "flashdreams-cosmos-predict2", - "flashdreams-self-forcing", - "flashdreams-t2v", -] - -[project.entry-points."flashdreams.runner_configs"] -t2v = "t2v_demo.runner:RUNNER_T2V" - -[project.entry-points."flashdreams.applications"] -t2v = "t2v_demo.app:create_app" - -[tool.uv.sources] -flashdreams = { workspace = true } -flashdreams-causal-forcing = { workspace = true } -flashdreams-cosmos-predict2 = { workspace = true } -flashdreams-self-forcing = { workspace = true } -flashdreams-t2v = { workspace = true } - -[tool.setuptools] -packages = ["t2v_demo"] -package-dir = { t2v_demo = "." } - -[tool.setuptools.package-data] -t2v_demo = ["web/*.js", "web/*.css"] diff --git a/apps/t2v_demo/runner.py b/apps/t2v_demo/runner.py deleted file mode 100644 index fd539f9fb..000000000 --- a/apps/t2v_demo/runner.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""``flashdreams-run t2v`` configuration and default replay runner.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Annotated, Any - -import tyro - -from flashdreams.infra.runner import Runner, RunnerConfig - -from .runtime import backend_choices, default_pipeline - - -@dataclass(kw_only=True) -class T2VDemoRunnerConfig(RunnerConfig): - """Configuration exposed by the ``flashdreams-run t2v`` slug.""" - - _target: type["T2VDemoRunner"] = field(default_factory=lambda: T2VDemoRunner) - launch_capability: Annotated[str | None, tyro.conf.Suppress] = ( - "t2v_demo.launch:LAUNCH_CAPABILITY" - ) - pipeline: Annotated[Any, tyro.conf.Suppress] = field( - default_factory=default_pipeline - ) - backend: str = "causal-forcing" - """Backend key: one of ``causal-forcing``, ``cosmos-predict2``, or ``self-forcing``.""" - - preset_id: str | None = None - prompt: str | None = None - total_blocks: int | None = None - pixel_height: int | None = None - pixel_width: int | None = None - fps: int | None = None - compile: bool | None = None - output: Path = Path("outputs/t2v.mp4") - - def __post_init__(self) -> None: - if self.backend not in backend_choices(): - raise ValueError( - f"Unknown T2V backend {self.backend!r}; choose one of " - f"{', '.join(backend_choices())}." - ) - - -class T2VDemoRunner(Runner[T2VDemoRunnerConfig, Any]): - """Default ``run`` mode delegates to the same app replay entry point.""" - - def __init__(self, config: T2VDemoRunnerConfig) -> None: - # The demo runtime owns pipeline construction so that replay and - # WebRTC share identical lifecycle handling. Avoid constructing a - # second Runner pipeline here. - self.config = config - - def run(self) -> None: - from .app import launch_t2v - - launch_t2v(config=self.config, mode="mp4") - - -RUNNER_T2V = T2VDemoRunnerConfig( - runner_name="t2v", - description="Text-to-video runtime demo (replay or WebRTC).", -) - -__all__ = ["RUNNER_T2V", "T2VDemoRunner", "T2VDemoRunnerConfig"] diff --git a/apps/t2v_demo/runtime.py b/apps/t2v_demo/runtime.py deleted file mode 100644 index 410ca7fb6..000000000 --- a/apps/t2v_demo/runtime.py +++ /dev/null @@ -1,226 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Compatibility bridge from ``flashdreams-run t2v`` to integration apps.""" - -from __future__ import annotations - -from dataclasses import dataclass, replace -from importlib import import_module -from importlib.metadata import entry_points -from typing import Any - -from t2v import ( - FIELD_FPS, - FIELD_PIXEL_HEIGHT, - FIELD_PIXEL_WIDTH, - FIELD_PROMPT, - FIELD_TOTAL_BLOCKS, - T2VDemoAdapter, - T2VInputProvider, - T2VModelConfig, - T2VRuntime, - T2VScenario, - T2VSession, - model_config_from_runner, -) - -from flashdreams.demo import Application, DemoAdapterApplication - - -@dataclass(frozen=True, slots=True) -class T2VBackendBridge: - """One legacy backend key routed to one integration-owned app.""" - - key: str - label: str - app_slug: str - app_module: str - config_module: str - - -_BACKENDS: dict[str, T2VBackendBridge] = { - "causal-forcing": T2VBackendBridge( - key="causal-forcing", - label="Causal-Forcing (Wan 2.1)", - app_slug="causal-forcing-t2v", - app_module="causal_forcing.t2v.app", - config_module="causal_forcing.config", - ), - "cosmos-predict2": T2VBackendBridge( - key="cosmos-predict2", - label="Cosmos Predict2", - app_slug="cosmos-predict2-t2v", - app_module="cosmos_predict2.t2v.app", - config_module="cosmos_predict2.config", - ), - "self-forcing": T2VBackendBridge( - key="self-forcing", - label="Self-Forcing (Wan 2.1)", - app_slug="self-forcing-t2v", - app_module="self_forcing.t2v.app", - config_module="self_forcing.config", - ), -} - - -def backend_choices() -> tuple[str, ...]: - """Return stable legacy CLI backend choices.""" - return tuple(_BACKENDS) - - -def backend_metadata() -> list[dict[str, Any]]: - """Return browser-safe backend metadata derived from integration configs.""" - metadata: list[dict[str, Any]] = [] - for bridge in _BACKENDS.values(): - model = model_from_backend(bridge.key) - metadata.append( - { - "key": bridge.key, - "label": bridge.label, - "default_preset": model.preset_id, - "presets": _t2v_preset_ids(bridge), - "application": bridge.app_slug, - } - ) - return metadata - - -def default_pipeline() -> Any: - """Return the legacy runner's suppressed default pipeline value.""" - return model_from_backend("causal-forcing").pipeline - - -def model_from_backend( - backend: str, - preset_id: str | None = None, -) -> T2VModelConfig: - """Resolve a legacy backend/preset selector to an integration-owned model.""" - bridge = _resolve_backend(backend) - if preset_id is None: - return _with_legacy_backend_option( - _default_model_from_application(bridge), bridge - ) - return _with_legacy_backend_option( - _model_from_runner_config(bridge, preset_id), bridge - ) - - -def make_adapter( - backend: str, - preset_id: str | None = None, - *, - write_download_artifact: bool = False, -) -> T2VDemoAdapter: - """Build an adapter from a legacy CLI/UI backend key.""" - return T2VDemoAdapter( - model=model_from_backend(backend, preset_id), - write_download_artifact=write_download_artifact, - ) - - -def _resolve_backend(value: str) -> T2VBackendBridge: - try: - return _BACKENDS[value] - except KeyError as exc: - raise ValueError( - f"Unknown backend {value!r}. Available backends: {', '.join(_BACKENDS)}." - ) from exc - - -def _default_model_from_application(bridge: T2VBackendBridge) -> T2VModelConfig: - app = _load_application(bridge) - if not isinstance(app, DemoAdapterApplication): - raise TypeError( - f"T2V application {bridge.app_slug!r} must return " - f"DemoAdapterApplication, got {type(app).__name__}." - ) - adapter = app.adapter - model = getattr(adapter, "model", None) - if not isinstance(model, T2VModelConfig): - raise TypeError(f"T2V application {bridge.app_slug!r} must use T2VDemoAdapter.") - return model - - -def _load_application(bridge: T2VBackendBridge) -> Application: - for entry_point in entry_points(group="flashdreams.applications"): - if entry_point.name == bridge.app_slug: - factory = entry_point.load() - return _coerce_application( - factory() if callable(factory) else factory, - app_slug=bridge.app_slug, - ) - factory = getattr(import_module(bridge.app_module), "create_app") - return _coerce_application(factory(), app_slug=bridge.app_slug) - - -def _coerce_application(value: object, *, app_slug: str) -> Application: - if not isinstance(value, Application): - raise TypeError( - f"T2V application {app_slug!r} must return Application, " - f"got {type(value).__name__}." - ) - return value - - -def _model_from_runner_config( - bridge: T2VBackendBridge, preset_id: str -) -> T2VModelConfig: - runner_configs = getattr(import_module(bridge.config_module), "RUNNER_CONFIGS") - try: - runner = runner_configs[preset_id] - except KeyError as exc: - raise ValueError( - f"Unknown {bridge.key} preset {preset_id!r}. Available presets: " - f"{', '.join(_t2v_preset_ids(bridge))}." - ) from exc - if not _is_t2v_preset(runner.runner_name): - raise ValueError( - f"Preset {preset_id!r} is not a T2V preset for backend {bridge.key!r}." - ) - default_model = _default_model_from_application(bridge) - return model_config_from_runner(model_id=default_model.model_id, runner=runner) - - -def _t2v_preset_ids(bridge: T2VBackendBridge) -> tuple[str, ...]: - runner_configs = getattr(import_module(bridge.config_module), "RUNNER_CONFIGS") - return tuple(name for name in runner_configs if _is_t2v_preset(name)) - - -def _is_t2v_preset(name: str) -> bool: - return "-t2v-" in name - - -def _with_legacy_backend_option( - model: T2VModelConfig, - bridge: T2VBackendBridge, -) -> T2VModelConfig: - return replace( - model, - runtime_options={ - **model.runtime_options, - "backend": bridge.key, - "application": bridge.app_slug, - }, - ) - - -__all__ = [ - "FIELD_FPS", - "FIELD_PIXEL_HEIGHT", - "FIELD_PIXEL_WIDTH", - "FIELD_PROMPT", - "FIELD_TOTAL_BLOCKS", - "T2VDemoAdapter", - "T2VBackendBridge", - "T2VInputProvider", - "T2VModelConfig", - "T2VRuntime", - "T2VScenario", - "T2VSession", - "backend_choices", - "backend_metadata", - "default_pipeline", - "make_adapter", - "model_from_backend", -] diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py deleted file mode 100644 index 295ff50de..000000000 --- a/apps/t2v_demo/tests/test_runner.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from pathlib import Path - -import pytest -import t2v.t2v as t2v_shell -import tomli -from t2v_demo import app -from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig -from t2v_demo.runtime import backend_metadata, model_from_backend - -from flashdreams.demo import ( - Application, - DemoAdapterApplication, -) -from flashdreams.runtime.demo import Mp4OutputSpec, RunResult - -pytestmark = pytest.mark.ci_cpu - - -def test_t2v_runner_slug_has_launch_capability() -> None: - assert RUNNER_T2V.runner_name == "t2v" - assert RUNNER_T2V.launch_capability == "t2v_demo.launch:LAUNCH_CAPABILITY" - - -def test_t2v_registers_application_entry_point() -> None: - pyproject_path = Path(__file__).parents[1] / "pyproject.toml" - pyproject = tomli.loads(pyproject_path.read_text()) - - assert pyproject["project"]["entry-points"]["flashdreams.applications"] == { - "t2v": "t2v_demo.app:create_app" - } - - -def test_t2v_demo_no_longer_owns_backend_presets() -> None: - package_dir = Path(__file__).parents[1] - - assert not (package_dir / "backends.py").exists() - assert not (package_dir / "presets.py").exists() - - -def test_t2v_create_app_exposes_public_application() -> None: - public_app = app.create_app( - T2VDemoRunnerConfig( - runner_name="t2v-test", - description="test", - backend="self-forcing", - prompt="A waterfall", - total_blocks=3, - ) - ) - - assert app.createApp is app.create_app - assert isinstance(public_app, Application) - assert isinstance(public_app, DemoAdapterApplication) - spec = public_app.spec - assert spec.model_id == "self-forcing-t2v" - assert spec.config is not None - assert spec.config.runtime_options["backend"] == "self-forcing" - assert spec.config.runtime_options["application"] == "self-forcing-t2v" - scenario = spec.scenario - assert isinstance(scenario, dict) - assert scenario["prompt"] == "A waterfall" - assert scenario["total_blocks"] == 3 - - -def test_t2v_backend_bridge_builds_neutral_model_config() -> None: - model = model_from_backend("self-forcing") - - assert isinstance(model, t2v_shell.T2VModelConfig) - assert model.model_id == "self-forcing-t2v" - assert model.runtime_options["backend"] == "self-forcing" - - -def test_t2v_backend_bridge_supports_integration_owned_preset() -> None: - model = model_from_backend( - "self-forcing", - "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope", - ) - - assert model.model_id == "self-forcing-t2v" - assert model.preset_id == "self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope" - assert model.total_blocks == 80 - - -def test_t2v_backend_metadata_is_derived_from_integrations() -> None: - metadata = {item["key"]: item for item in backend_metadata()} - - assert metadata["self-forcing"]["default_preset"] == ( - "self-forcing-wan2.1-t2v-1.3b" - ) - assert "self-forcing-wan2.1-t2v-1.3b-taehv" in metadata["self-forcing"]["presets"] - assert metadata["cosmos-predict2"]["application"] == "cosmos-predict2-t2v" - - -def test_runner_mp4_launch_uses_demo_entrypoint( - monkeypatch: pytest.MonkeyPatch, -) -> None: - captured: dict[str, object] = {} - - def fake_run_application_replay(*, app: Application) -> RunResult: - captured["app"] = app - return RunResult(status="completed") - - monkeypatch.setattr(app, "run_application_replay", fake_run_application_replay) - config = T2VDemoRunnerConfig( - runner_name="t2v", - description="test", - backend="self-forcing", - prompt="A waterfall", - total_blocks=3, - ) - - app.launch_t2v( - config=config, - mode="mp4", - output_overrides={"path": "outputs/test.mp4", "fps": 24}, - ) - - public_app = captured["app"] - assert isinstance(public_app, DemoAdapterApplication) - scenario = public_app.spec.scenario - assert isinstance(scenario, dict) - assert scenario["prompt"] == "A waterfall" - assert scenario["total_blocks"] == 3 - output = public_app.spec.output - assert isinstance(output, Mp4OutputSpec) - assert str(output.path) == "outputs/test.mp4" - assert output.fps == 24 - assert output.output_layout == "tchw" diff --git a/apps/t2v_demo/web/adapter.css b/apps/t2v_demo/web/adapter.css deleted file mode 100644 index 66e816f29..000000000 --- a/apps/t2v_demo/web/adapter.css +++ /dev/null @@ -1,3 +0,0 @@ -.t2vPanel { display: grid; gap: .75rem; padding: 1rem; max-width: 32rem; } -.t2vPanel textarea { display: block; width: 100%; margin-top: .35rem; resize: vertical; } -.t2vPanel button { margin-right: .5rem; } diff --git a/apps/t2v_demo/web/adapter.js b/apps/t2v_demo/web/adapter.js deleted file mode 100644 index 7c56a622a..000000000 --- a/apps/t2v_demo/web/adapter.js +++ /dev/null @@ -1,23 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -/** Model metadata only; shared WebRTC UI renders prompt/video controls. */ -export default { - modelName: "Text-to-Video", - async mount(context) { - const response = await fetch("/api/t2v/config") - if (!response.ok) return - const config = await response.json() - const selected = config.backends?.find((backend) => backend.key === config.selected_backend) - context.setModelName(selected?.label || config.selected_backend || "Text-to-Video") - }, - promptGeneration: { - endpoint: "/api/t2v/prompt", - label: "Describe the video", - placeholder: "A cinematic drone shot over snowy mountains at sunrise", - generateLabel: "Generate video", - downloadEndpoint: "/api/t2v/download", - playbackEndpoint: "/api/t2v/playback", - hideControls: true, - }, -} diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index 37c83ebf2..1648bfd94 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -6,6 +6,7 @@ from __future__ import annotations import argparse +import asyncio import sys from collections.abc import Callable, Sequence from dataclasses import dataclass, replace @@ -193,29 +194,42 @@ def run_application_webrtc( adapter = app.adapter scenario = adapter.prepare_scenario(spec) runtime = adapter.create_runtime(web_config) - manager = _create_application_webrtc_manager( - runtime=runtime, - output=output, - spec=spec, - scenario=scenario, - adapter=adapter, - ) - app_resources = _create_application_webrtc_resources( - adapter=adapter, - manager=manager, - output=output, - spec=spec, - ) + manager: Any | None = None + try: + manager = _create_application_webrtc_manager( + runtime=runtime, + output=output, + spec=spec, + scenario=scenario, + adapter=adapter, + ) + app_resources = _create_application_webrtc_resources( + adapter=adapter, + manager=manager, + output=output, + spec=spec, + ) - from flashdreams.serving.webrtc.demo import serve_webrtc_demo + from flashdreams.serving.webrtc.demo import serve_webrtc_demo - return serve_webrtc_demo( - output=output, - model_id=spec.model_id, - session_manager=manager, - app_resources=app_resources, - world_rank=context.world_rank, - ) + return serve_webrtc_demo( + output=output, + model_id=spec.model_id, + session_manager=manager, + app_resources=app_resources, + world_rank=context.world_rank, + ) + except BaseException as exc: + # Once the runtime exists, aiohttp shutdown is not guaranteed to run + # until the server fully starts. Clean up through the same manager/host + # ownership path that normal WebRTC shutdown uses, while preserving the + # startup failure as the primary exception. + _cleanup_application_webrtc_startup_failure( + manager=manager, + runtime=runtime, + primary_error=exc, + ) + raise def create_demo_application( @@ -270,6 +284,23 @@ def _create_application_webrtc_manager( ) +def _cleanup_application_webrtc_startup_failure( + *, + manager: Any | None, + runtime: Any, + primary_error: BaseException, +) -> None: + try: + if manager is None: + runtime.close() + return + asyncio.run(manager.shutdown()) + except BaseException as cleanup_error: + add_note = getattr(primary_error, "add_note", None) + if callable(add_note): + add_note(f"Additional WebRTC startup cleanup error: {cleanup_error!r}") + + def _create_application_webrtc_resources( *, adapter: DemoAdapter, diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index 0e54e3c86..a85b1df23 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -9,10 +9,12 @@ from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path +from types import SimpleNamespace from typing import Any import pytest +import flashdreams.demo.app as demo_app import flashdreams.demo.runner as demo_runner from flashdreams.demo import ( Application, @@ -78,6 +80,7 @@ SingleSessionAdmissionPolicy, StepPipeline, UserInputWindow, + WebRTCOutputSpec, ) from flashdreams.runtime.output import OutputArtifact from flashdreams.runtime.types import StepRequirements @@ -876,6 +879,40 @@ def test_run_application_replay_uses_demo_adapter_spec() -> None: assert runtime.closed +def test_run_application_webrtc_closes_runtime_when_server_startup_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter = _FakeDemoAdapter() + app = DemoAdapterApplication( + adapter=adapter, + spec=DemoSpec( + model_id="fake-demo", + input_mode="webrtc", + output=WebRTCOutputSpec(), + config=InferenceConfig(model_id="fake-demo", device="cuda:7"), + ), + ) + + monkeypatch.setattr( + demo_app, + "initialize_cuda_distributed", + lambda **_: SimpleNamespace(device="cuda:0", world_rank=0), + ) + + import flashdreams.serving.webrtc.demo as webrtc_demo + + def fail_to_serve(**_: object) -> object: + raise RuntimeError("server startup failed") + + monkeypatch.setattr(webrtc_demo, "serve_webrtc_demo", fail_to_serve) + + with pytest.raises(RuntimeError, match="server startup failed"): + run_application_webrtc(app=app) + + assert len(adapter.runtimes) == 1 + assert adapter.runtimes[0].closed + + def test_demo_application_can_be_built_from_callbacks() -> None: adapter = _FakeDemoAdapter() app = create_demo_application( diff --git a/uv.lock b/uv.lock index 5065e7bc7..a1fddd90d 100644 --- a/uv.lock +++ b/uv.lock @@ -30,7 +30,6 @@ members = [ "flashdreams-sana-wm", "flashdreams-self-forcing", "flashdreams-t2v", - "flashdreams-t2v-demo", "flashdreams-wan21", "flashdreams-wan22", "ludus-renderer", @@ -1436,27 +1435,6 @@ dependencies = [ [package.metadata] requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] -[[package]] -name = "flashdreams-t2v-demo" -version = "0.1.0" -source = { editable = "apps/t2v_demo" } -dependencies = [ - { name = "flashdreams", extra = ["serving"] }, - { name = "flashdreams-causal-forcing" }, - { name = "flashdreams-cosmos-predict2" }, - { name = "flashdreams-self-forcing" }, - { name = "flashdreams-t2v" }, -] - -[package.metadata] -requires-dist = [ - { name = "flashdreams", extras = ["serving"], editable = "flashdreams" }, - { name = "flashdreams-causal-forcing", editable = "integrations/causal_forcing" }, - { name = "flashdreams-cosmos-predict2", editable = "integrations/cosmos_predict2" }, - { name = "flashdreams-self-forcing", editable = "integrations/self_forcing" }, - { name = "flashdreams-t2v", editable = "apps/t2v" }, -] - [[package]] name = "flashdreams-wan21" version = "0.1.0" From 83449e687b62fd888bcd8c167276c5551d3789cb Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 17:37:17 +0000 Subject: [PATCH 28/33] Add wan21 and wan22 T2V apps with usage README Register `wan21-t2v` and `fastvideo-causal-wan22-t2v` under the `flashdreams.applications` entry-point group so both models launch directly from `flashdreams-run` alongside the existing three T2V apps. Each integration gains a `t2v/` subpackage wrapping its existing runner config, plus the `flashdreams-t2v` workspace dependency. No shell changes were needed: wan22 is autoregressive like the existing streaming apps, and wan21 T2V is single-step, which `model_config_from_runner` already covers via its `total_blocks` fallback of 1. Restore the T2V README, lost with `apps/t2v_demo`, documenting the five slugs, the three launch modes, and the scenario and output overrides. --- apps/t2v/README.md | 92 +++++++++++++++++++ .../fastvideo_causal_wan22/t2v/__init__.py | 8 ++ .../fastvideo_causal_wan22/t2v/app.py | 24 +++++ .../fastvideo_causal_wan22/pyproject.toml | 5 + integrations/wan21/pyproject.toml | 5 + integrations/wan21/wan21/t2v/__init__.py | 8 ++ integrations/wan21/wan21/t2v/app.py | 24 +++++ uv.lock | 4 + 8 files changed, 170 insertions(+) create mode 100644 apps/t2v/README.md create mode 100644 integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/__init__.py create mode 100644 integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/app.py create mode 100644 integrations/wan21/wan21/t2v/__init__.py create mode 100644 integrations/wan21/wan21/t2v/app.py diff --git a/apps/t2v/README.md b/apps/t2v/README.md new file mode 100644 index 000000000..66bfacae7 --- /dev/null +++ b/apps/t2v/README.md @@ -0,0 +1,92 @@ +# FlashDreams text-to-video app + +A model-neutral text-to-video shell built on the shared +`flashdreams.runtime.demo` session lifecycle. Each integration owns its pipeline +and checkpoint configuration and registers a slug under the +`flashdreams.applications` entry-point group; this app owns only prompt input, +output selection, and the thin runtime adapter. + +## Available applications + +Every slug below launches through `flashdreams-run` and carries its model's +defaults, so a bare invocation needs no flags. + +| Slug | Preset | Blocks | Size | FPS | +| --- | --- | --- | --- | --- | +| `causal-forcing-t2v` | `causal-forcing-wan2.1-t2v-1.3b-chunkwise` | 60 | 832x480 | 16 | +| `cosmos-predict2-t2v` | `cosmos2-t2v-2b-720p` | 1 | 1280x720 | 16 | +| `fastvideo-causal-wan22-t2v` | `fastvideo-causal-wan2.2-t2v-14b` | 60 | 832x480 | 16 | +| `self-forcing-t2v` | `self-forcing-wan2.1-t2v-1.3b` | 60 | 832x480 | 16 | +| `wan21-t2v` | `wan21-t2v-1.3b-480p` | 1 | 832x480 | 16 | + +Models with one block generate a single chunk; the rest roll out +autoregressively for `total_blocks` chunks. + +## Running + +The launch mode is positional: `flashdreams-run `, where `` +is `mp4`, `null`, or `webrtc`. + +Write an MP4, defaulting to `outputs/.mp4`: + +```bash +uv run flashdreams-run wan21-t2v mp4 +``` + +Serve streamed WebRTC output, defaulting to `127.0.0.1:8080`: + +```bash +uv run flashdreams-run self-forcing-t2v webrtc +``` + +The browser UI accepts a prompt before opening a generation session, plays +emitted chunks as they arrive, and records the received stream for download. + +Generate without writing anything, which is the cheapest way to exercise a +model end to end: + +```bash +uv run flashdreams-run causal-forcing-t2v null +``` + +## Overrides + +`--scenario.KEY VALUE` overrides generation settings: `prompt`, `total_blocks`, +`pixel_height`, `pixel_width`, and `fps`. + +```bash +uv run flashdreams-run fastvideo-causal-wan22-t2v mp4 \ + --scenario.prompt "a red fox padding through fresh snow" \ + --scenario.total_blocks 8 +``` + +`--output.KEY VALUE` overrides sink settings, which vary by mode. MP4 accepts +`path`, `fps`, `output_layout`, and `move_to_cpu`; `null` accepts +`store_results`; WebRTC accepts `host`, `port`, `video_width`, `video_height`, +`warmup_chunks`, and `client_liveness_timeout_s`, among others. WebRTC also +takes `--host` and `--port` directly. + +```bash +uv run flashdreams-run wan21-t2v mp4 --output.path outputs/fox.mp4 +uv run flashdreams-run self-forcing-t2v webrtc --host 0.0.0.0 --port 8099 +``` + +Append `--no-instantiate` to resolve and print a launch without loading the +model, which validates a command in seconds: + +```bash +uv run flashdreams-run wan21-t2v mp4 --scenario.total_blocks 2 --no-instantiate +``` + +## Image-to-video + +The app shell conditions on a prompt only, so the image-to-video presets are not +available as application slugs. Run them through their runner slugs instead, +which accept `--prompt` and `--image-path`: + +```bash +uv run flashdreams-run wan21-i2v-14b-480p --image-path first_frame.png +``` + +The same applies to `causal-forcing-wan2.1-i2v-1.3b-framewise` and +`cosmos2-i2v-2b-720p`. diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/__init__.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/__init__.py new file mode 100644 index 000000000..6fe47efc9 --- /dev/null +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FastVideo CausalWan 2.2 T2V public demo app.""" + +from fastvideo_causal_wan22.t2v.app import MODEL, create_app, createApp + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/app.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/app.py new file mode 100644 index 000000000..ea9b18298 --- /dev/null +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/t2v/app.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public T2V app entry for the default FastVideo CausalWan 2.2 model.""" + +from t2v import create_t2v_application, model_config_from_runner + +from fastvideo_causal_wan22.config import RUNNER_WAN22_T2V_14B +from flashdreams.demo import Application + +MODEL = model_config_from_runner( + model_id="fastvideo-causal-wan22-t2v", + runner=RUNNER_WAN22_T2V_14B, +) + + +def create_app() -> Application: + """Create the default FastVideo CausalWan 2.2 T2V application.""" + return create_t2v_application(model=MODEL) + + +createApp = create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/fastvideo_causal_wan22/pyproject.toml b/integrations/fastvideo_causal_wan22/pyproject.toml index a3dd56636..d72a02f20 100644 --- a/integrations/fastvideo_causal_wan22/pyproject.toml +++ b/integrations/fastvideo_causal_wan22/pyproject.toml @@ -25,11 +25,13 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "flashdreams", + "flashdreams-t2v", "mediapy>=1.1", ] [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [project.optional-dependencies] dev = [ @@ -43,6 +45,9 @@ dev = [ [project.entry-points."flashdreams.runner_configs"] "fastvideo-causal-wan2.2-t2v-14b" = "fastvideo_causal_wan22.config:RUNNER_WAN22_T2V_14B" +[project.entry-points."flashdreams.applications"] +fastvideo-causal-wan22-t2v = "fastvideo_causal_wan22.t2v.app:create_app" + [tool.setuptools.packages.find] include = ["fastvideo_causal_wan22*"] exclude = ["tests"] diff --git a/integrations/wan21/pyproject.toml b/integrations/wan21/pyproject.toml index 334a302a9..386512926 100644 --- a/integrations/wan21/pyproject.toml +++ b/integrations/wan21/pyproject.toml @@ -25,12 +25,14 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "flashdreams", + "flashdreams-t2v", "mediapy>=1.1", "opencv-python-headless>=4.5", ] [tool.uv.sources] flashdreams = { workspace = true } +flashdreams-t2v = { workspace = true } [project.optional-dependencies] dev = [ @@ -45,6 +47,9 @@ dev = [ "wan21-t2v-1.3b-480p" = "wan21.config:RUNNER_WAN21_T2V_1PT3B_480P" "wan21-i2v-14b-480p" = "wan21.config:RUNNER_WAN21_I2V_14B_480P" +[project.entry-points."flashdreams.applications"] +wan21-t2v = "wan21.t2v.app:create_app" + [tool.setuptools.packages.find] include = ["wan21*"] exclude = ["tests"] diff --git a/integrations/wan21/wan21/t2v/__init__.py b/integrations/wan21/wan21/t2v/__init__.py new file mode 100644 index 000000000..c0f7f1dc0 --- /dev/null +++ b/integrations/wan21/wan21/t2v/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Wan 2.1 T2V public demo app.""" + +from wan21.t2v.app import MODEL, create_app, createApp + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/integrations/wan21/wan21/t2v/app.py b/integrations/wan21/wan21/t2v/app.py new file mode 100644 index 000000000..0c7228d6d --- /dev/null +++ b/integrations/wan21/wan21/t2v/app.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public T2V app entry for the default Wan 2.1 model.""" + +from t2v import create_t2v_application, model_config_from_runner + +from flashdreams.demo import Application +from wan21.config import RUNNER_WAN21_T2V_1PT3B_480P + +MODEL = model_config_from_runner( + model_id="wan21-t2v", + runner=RUNNER_WAN21_T2V_1PT3B_480P, +) + + +def create_app() -> Application: + """Create the default Wan 2.1 T2V application.""" + return create_t2v_application(model=MODEL) + + +createApp = create_app + +__all__ = ["MODEL", "createApp", "create_app"] diff --git a/uv.lock b/uv.lock index a1fddd90d..32552fd6b 100644 --- a/uv.lock +++ b/uv.lock @@ -1163,6 +1163,7 @@ version = "0.1.0" source = { editable = "integrations/fastvideo_causal_wan22" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1174,6 +1175,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, { name = "mediapy", specifier = ">=1.1" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, ] @@ -1441,6 +1443,7 @@ version = "0.1.0" source = { editable = "integrations/wan21" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, { name = "opencv-python-headless" }, ] @@ -1453,6 +1456,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, { name = "mediapy", specifier = ">=1.1" }, { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, From a1c07641beb5a7c81ba398a63eafd3314e4e281a Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 17:54:33 +0000 Subject: [PATCH 29/33] Close apps when external runtime host is already closed Ensure public runner cleanup still invokes Application.close() when an externally supplied RuntimeHost has already stopped accepting dispatched work. Add sync and async regression coverage for the closed-host fallback so application-owned runtimes are not orphaned. --- apps/t2v/README.md | 14 ++ flashdreams/flashdreams/demo/runner.py | 34 ++++- .../tests/test_demo_application_api.py | 122 +++++++++++++++++- 3 files changed, 164 insertions(+), 6 deletions(-) diff --git a/apps/t2v/README.md b/apps/t2v/README.md index 66bfacae7..625b1084e 100644 --- a/apps/t2v/README.md +++ b/apps/t2v/README.md @@ -6,6 +6,20 @@ and checkpoint configuration and registers a slug under the `flashdreams.applications` entry-point group; this app owns only prompt input, output selection, and the thin runtime adapter. +## Setup + +Install the workspace from the repo root, which registers every application +slug below, then set a token for checkpoint downloads: + +```bash +uv sync --extra runners +export HF_TOKEN= +``` + +Developers and maintainers should use `uv sync --extra dev --extra runners` +instead. To build only one model's dependencies rather than the whole workspace, +see that model's page under `docs/source/models/`. + ## Available applications Every slug below launches through `flashdreams-run` and carries its model's diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index d49c36fd2..992cf8b22 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -534,7 +534,11 @@ def close_app() -> None: if cleanup.closed: return if not invoked and host.is_closed: - errors.append(_closed_host_cleanup_error(exc)) + _close_application_after_closed_host( + errors=errors, + cleanup=cleanup, + host_error=exc, + ) return errors.append(exc) @@ -560,11 +564,33 @@ def close_app() -> None: if cleanup.closed: return if not invoked and host.is_closed: - errors.append(_closed_host_cleanup_error(exc)) + _close_application_after_closed_host( + errors=errors, + cleanup=cleanup, + host_error=exc, + ) return errors.append(exc) +def _close_application_after_closed_host( + *, + errors: list[Exception], + cleanup: _ApplicationCleanup, + host_error: Exception, +) -> None: + # An externally owned host may already be torn down by the time runner + # cleanup runs. At that point worker dispatch is impossible, so the + # idempotent app cleanup is the last leak-prevention fallback. + fallback_errors: list[Exception] = [] + _record_cleanup_error(fallback_errors, cleanup.close) + if cleanup.closed: + return + cleanup_error = _closed_host_cleanup_error(host_error) + _record_cleanup_notes(cleanup_error, fallback_errors) + errors.append(cleanup_error) + + def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: if not errors: return @@ -617,8 +643,8 @@ def _has_primary_outcome( def _closed_host_cleanup_error(exc: Exception) -> RuntimeError: try: raise RuntimeError( - "Application cleanup skipped because the RuntimeHost is closed; " - "app.close() must run on the model worker." + "Application cleanup could not be dispatched because the RuntimeHost " + "is closed." ) from exc except RuntimeError as cleanup_error: return cleanup_error diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index a85b1df23..c9cc229a3 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -386,7 +386,7 @@ def test_runner_external_host_close_hook_closes_public_app() -> None: assert app.closed -def test_runner_skips_direct_app_close_for_closed_host() -> None: +def test_runner_does_not_init_public_app_when_host_starts_closed() -> None: app = _RunnerFakeApplication(total_steps=1) host = RuntimeHost(_ExternalApplicationRuntime(app)) host.close() @@ -405,6 +405,30 @@ def test_runner_skips_direct_app_close_for_closed_host() -> None: assert not app.closed +def test_runner_closes_public_app_when_external_host_disappears() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) + + try: + result = Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="external-host-disappears", + driver=_CloseHostWithoutHooksDriver(), + ), + model_id="fake-runner", + ).run() + finally: + host.close() + + assert result.status == "completed" + assert host.is_closed + assert app.closed + + def test_runner_closes_public_app_when_context_cleanup_fails() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -514,7 +538,9 @@ async def test_runner_run_async_delegates_to_async_session_helper() -> None: @pytest.mark.asyncio -async def test_runner_run_async_skips_direct_app_close_for_closed_host() -> None: +async def test_runner_run_async_does_not_init_public_app_when_host_starts_closed() -> ( + None +): app = _RunnerFakeApplication(total_steps=1) host = RuntimeHost(_ExternalApplicationRuntime(app)) host.close() @@ -534,6 +560,33 @@ async def test_runner_run_async_skips_direct_app_close_for_closed_host() -> None assert not app.closed +@pytest.mark.asyncio +async def test_runner_run_async_closes_public_app_when_external_host_disappears() -> ( + None +): + app = _RunnerFakeApplication(total_steps=1) + host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) + + try: + result = await Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="async-external-host-disappears", + driver=_AsyncCloseHostWithoutHooksDriver(), + ), + model_id="fake-runner", + ).run_async() + finally: + host.close() + + assert result.status == "completed" + assert host.is_closed + assert app.closed + + @pytest.mark.asyncio async def test_runner_run_async_external_host_close_hook_closes_public_app() -> None: app = _RunnerFakeApplication(total_steps=1) @@ -1424,6 +1477,39 @@ def close(self) -> None: self.closed = True +class _ExternallyClosedWithoutHooksRuntimeHost(RuntimeHost): + def __init__(self, runtime: _ExternalApplicationRuntime) -> None: + super().__init__(runtime) + self._externally_closed = False + + @property + def is_healthy(self) -> bool: + return super().is_healthy and not self._externally_closed + + @property + def is_closed(self) -> bool: + return self._externally_closed or super().is_closed + + def close_without_hooks(self) -> None: + self._externally_closed = True + + def call(self, func: Any, /, *args: object, **kwargs: object) -> Any: + if self._externally_closed: + raise RuntimeError("runtime host is closed") + return super().call(func, *args, **kwargs) + + async def call_async( + self, + func: Any, + /, + *args: object, + **kwargs: object, + ) -> Any: + if self._externally_closed: + raise RuntimeError("runtime host is closed") + return await super().call_async(func, *args, **kwargs) + + class _FailingCloseMetricsRecorder(InMemorySessionMetricsRecorder): def close(self) -> Any: raise RuntimeError("run metrics close failed") @@ -1566,6 +1652,22 @@ def run_one_session( return session_edges.close_result(status="completed") +class _CloseHostWithoutHooksDriver: + def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + assert isinstance(host, _ExternallyClosedWithoutHooksRuntimeHost) + host.call(provider.close) + host.close_without_hooks() + return session_edges.close_result(status="completed") + + @dataclass(slots=True) class _AsyncRecordingRunMode: io_handler: _RecordingIOHandler @@ -1680,6 +1782,22 @@ async def run_one_session( return session_edges.close_result(status="completed") +class _AsyncCloseHostWithoutHooksDriver: + async def run_one_session( + self, + *, + host: RuntimeHost, + provider: Any, + session_edges: SessionEdges, + pipeline: StepPipeline, + ) -> RunResult: + del pipeline + assert isinstance(host, _ExternallyClosedWithoutHooksRuntimeHost) + await host.call_async(provider.close) + host.close_without_hooks() + return session_edges.close_result(status="completed") + + class _AsyncRunModeInputSource: is_finite = False is_deterministic = False From 2e13082073ad0bc22f320953c38caec91208460a Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 18:11:33 +0000 Subject: [PATCH 30/33] Harden WebRTC manager shutdown cleanup Make BaseWebRTCSessionManager.shutdown attempt every owned cleanup step even when active-session shutdown fails. Preserve the first cleanup error while recording later failures as notes, and add regression coverage that shared RuntimeHost cleanup still runs after session and context close failures. --- .../flashdreams/serving/webrtc/manager.py | 49 ++++++++++++++--- flashdreams/tests/test_webrtc_manager.py | 54 ++++++++++++++++++- 2 files changed, 95 insertions(+), 8 deletions(-) diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 44190e592..733eaeebc 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -658,6 +658,31 @@ async def close(self) -> None: await self.peer_connection.close() +async def _record_shutdown_cleanup_error( + errors: list[BaseException], + cleanup: Callable[..., Any], + /, + *args: Any, +) -> None: + try: + result = cleanup(*args) + if inspect.isawaitable(result): + await result + except BaseException as exc: + errors.append(exc) + + +def _raise_first_shutdown_cleanup_error(errors: list[BaseException]) -> None: + if not errors: + return + first = errors[0] + add_note = getattr(first, "add_note", None) + for extra in errors[1:]: + if callable(add_note): + add_note(f"Additional WebRTC shutdown cleanup error: {extra!r}") + raise first + + class BaseWebRTCSessionManager(Generic[_RuntimeT, _RuntimeConfigT]): """Owns one active WebRTC session and forwards actions into a model runtime.""" @@ -1641,25 +1666,35 @@ async def _client_liveness_watchdog( raise async def shutdown(self) -> None: - await self.close_active_session() + errors: list[BaseException] = [] + await _record_shutdown_cleanup_error(errors, self.close_active_session) if self._shared_context is not None: - await self._shared_context.close_async() + await _record_shutdown_cleanup_error( + errors, + self._shared_context.close_async, + ) if self._shared_host is not None: - await asyncio.to_thread(self._shared_host.close) + await _record_shutdown_cleanup_error( + errors, + asyncio.to_thread, + self._shared_host.close, + ) self._shared_context = None self._shared_host = None self._shared_runtime_adapter = None if self._shared_video_encoder is not None: - self._shared_video_encoder.close() + await _record_shutdown_cleanup_error( + errors, + self._shared_video_encoder.close, + ) self._shared_video_encoder = None if not self._owns_shared_host: close = getattr(self._runtime, "close", None) if callable(close): - result = close() - if inspect.isawaitable(result): - await result + await _record_shutdown_cleanup_error(errors, close) self._runtime_ready = False self._warmup_complete = False + _raise_first_shutdown_cleanup_error(errors) def wait_for_termination(self) -> None: wait = getattr(self._runtime, "wait_for_termination", None) diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index bafe2b1a8..cb8589d49 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -13,13 +13,14 @@ from flashdreams.runtime import ( InferenceInput, + InferenceSession, StepRequest, StepRequirements, StepResult, UserInputEvent, UserInputs, ) -from flashdreams.runtime.demo import RunResult +from flashdreams.runtime.demo import RunResult, RuntimeHost from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc import manager as manager_module from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult @@ -292,6 +293,57 @@ def _managed_session( return managed, video_track, peer, channel +@pytest.mark.asyncio +async def test_shutdown_closes_owned_shared_host_when_session_close_fails() -> None: + class _ClosingRuntime: + def __init__(self) -> None: + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + raise AssertionError("shutdown test must not start sessions") + + def close(self) -> None: + self.closed = True + + class _FailingVideoTrack(_FakeVideoTrack): + async def close(self) -> None: + raise RuntimeError("active session close failed") + + class _FailingSharedContext: + def __init__(self) -> None: + self.closed = False + + async def close_async(self) -> None: + self.closed = True + raise RuntimeError("context close failed") + + runtime = _ClosingRuntime() + host = RuntimeHost(runtime) + manager = _make_manager(_BaseTestManager, runtime, shared_host=host) + context = _FailingSharedContext() + managed, _video_track, _peer, _channel = _managed_session(runtime) + managed.video_track = _FailingVideoTrack() # ty:ignore[invalid-assignment] + manager._active_session = managed + manager._shared_context = cast(Any, context) + manager._runtime_ready = True + manager._warmup_complete = True + + with pytest.raises(RuntimeError, match="active session close failed") as exc_info: + await manager.shutdown() + + assert context.closed + assert host.is_closed + assert runtime.closed + assert manager._shared_context is None + assert manager._shared_host is None + assert manager._shared_runtime_adapter is None + assert not manager._runtime_ready + assert not manager._warmup_complete + notes = getattr(exc_info.value, "__notes__", ()) + assert any("context close failed" in note for note in notes) + + def test_record_user_event_rejects_full_queue( monkeypatch: pytest.MonkeyPatch, ) -> None: From 36a2a5c74656d53b47b5fdfcebd0df6896e6122b Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 18:31:40 +0000 Subject: [PATCH 31/33] Avoid off-worker app cleanup after closed host Do not call Application.close() directly when an external RuntimeHost is already closed before runner cleanup can dispatch onto the model worker. Report the invalid lifecycle instead, preserving the worker-affinity boundary required for model/runtime cleanup. Update sync and async runner tests to cover the closed-host failure path. --- flashdreams/flashdreams/demo/runner.py | 33 ++-------- .../tests/test_demo_application_api.py | 62 ++++++++++--------- 2 files changed, 38 insertions(+), 57 deletions(-) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index 992cf8b22..d28d08fdc 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -534,11 +534,7 @@ def close_app() -> None: if cleanup.closed: return if not invoked and host.is_closed: - _close_application_after_closed_host( - errors=errors, - cleanup=cleanup, - host_error=exc, - ) + errors.append(_closed_host_cleanup_error(exc)) return errors.append(exc) @@ -564,33 +560,11 @@ def close_app() -> None: if cleanup.closed: return if not invoked and host.is_closed: - _close_application_after_closed_host( - errors=errors, - cleanup=cleanup, - host_error=exc, - ) + errors.append(_closed_host_cleanup_error(exc)) return errors.append(exc) -def _close_application_after_closed_host( - *, - errors: list[Exception], - cleanup: _ApplicationCleanup, - host_error: Exception, -) -> None: - # An externally owned host may already be torn down by the time runner - # cleanup runs. At that point worker dispatch is impossible, so the - # idempotent app cleanup is the last leak-prevention fallback. - fallback_errors: list[Exception] = [] - _record_cleanup_error(fallback_errors, cleanup.close) - if cleanup.closed: - return - cleanup_error = _closed_host_cleanup_error(host_error) - _record_cleanup_notes(cleanup_error, fallback_errors) - errors.append(cleanup_error) - - def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: if not errors: return @@ -644,7 +618,8 @@ def _closed_host_cleanup_error(exc: Exception) -> RuntimeError: try: raise RuntimeError( "Application cleanup could not be dispatched because the RuntimeHost " - "is closed." + "is closed. External RuntimeHost owners must keep the host open until " + "runner cleanup completes or close it through RuntimeHost.close()." ) from exc except RuntimeError as cleanup_error: return cleanup_error diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index c9cc229a3..e45768a62 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -405,28 +405,31 @@ def test_runner_does_not_init_public_app_when_host_starts_closed() -> None: assert not app.closed -def test_runner_closes_public_app_when_external_host_disappears() -> None: +def test_runner_raises_when_external_host_disappears_before_cleanup() -> None: app = _RunnerFakeApplication(total_steps=1) host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) try: - result = Runner( - io_handler=_RecordingIOHandler(), - app=app, - host=host, - run_mode=_AsyncRecordingRunMode( - _RecordingIOHandler(), - name="external-host-disappears", - driver=_CloseHostWithoutHooksDriver(), - ), - model_id="fake-runner", - ).run() + with pytest.raises( + RuntimeError, + match="Application cleanup could not be dispatched", + ): + Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="external-host-disappears", + driver=_CloseHostWithoutHooksDriver(), + ), + model_id="fake-runner", + ).run() finally: host.close() - assert result.status == "completed" assert host.is_closed - assert app.closed + assert not app.closed def test_runner_closes_public_app_when_context_cleanup_fails() -> None: @@ -561,30 +564,33 @@ async def test_runner_run_async_does_not_init_public_app_when_host_starts_closed @pytest.mark.asyncio -async def test_runner_run_async_closes_public_app_when_external_host_disappears() -> ( +async def test_runner_run_async_raises_when_external_host_disappears_before_cleanup() -> ( None ): app = _RunnerFakeApplication(total_steps=1) host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) try: - result = await Runner( - io_handler=_RecordingIOHandler(), - app=app, - host=host, - run_mode=_AsyncRecordingRunMode( - _RecordingIOHandler(), - name="async-external-host-disappears", - driver=_AsyncCloseHostWithoutHooksDriver(), - ), - model_id="fake-runner", - ).run_async() + with pytest.raises( + RuntimeError, + match="Application cleanup could not be dispatched", + ): + await Runner( + io_handler=_RecordingIOHandler(), + app=app, + host=host, + run_mode=_AsyncRecordingRunMode( + _RecordingIOHandler(), + name="async-external-host-disappears", + driver=_AsyncCloseHostWithoutHooksDriver(), + ), + model_id="fake-runner", + ).run_async() finally: host.close() - assert result.status == "completed" assert host.is_closed - assert app.closed + assert not app.closed @pytest.mark.asyncio From 5b11641e25bb7288eb4aaa3a13644a74e8a2ed97 Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 18:52:46 +0000 Subject: [PATCH 32/33] Document external host cleanup contract --- flashdreams/flashdreams/demo/runner.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/flashdreams/flashdreams/demo/runner.py b/flashdreams/flashdreams/demo/runner.py index d28d08fdc..172875f05 100644 --- a/flashdreams/flashdreams/demo/runner.py +++ b/flashdreams/flashdreams/demo/runner.py @@ -615,6 +615,10 @@ def _has_primary_outcome( def _closed_host_cleanup_error(exc: Exception) -> RuntimeError: + # Do not call Application.close() directly here. Application cleanup can own + # model or CUDA state, so it must be dispatched through RuntimeHost's worker. + # External host owners are expected to keep the host open through runner + # teardown or close it via RuntimeHost.close(), which runs registered hooks. try: raise RuntimeError( "Application cleanup could not be dispatched because the RuntimeHost " From cd3e3617ca96103036cdcf10ab12f22ad020abce Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Thu, 13 Aug 2026 19:16:58 +0000 Subject: [PATCH 33/33] Clean up WebRTC distributed state on startup failure --- flashdreams/flashdreams/demo/app.py | 63 ++++++++- .../flashdreams/runtime/demo/bootstrap.py | 37 ++++++ .../flashdreams/serving/webrtc/bootstrap.py | 122 ++++++++++++++---- .../tests/test_demo_application_api.py | 12 ++ flashdreams/tests/test_webrtc_bootstrap.py | 112 ++++++++++++++++ 5 files changed, 316 insertions(+), 30 deletions(-) diff --git a/flashdreams/flashdreams/demo/app.py b/flashdreams/flashdreams/demo/app.py index 1648bfd94..a8d720a9e 100644 --- a/flashdreams/flashdreams/demo/app.py +++ b/flashdreams/flashdreams/demo/app.py @@ -7,6 +7,7 @@ import argparse import asyncio +import inspect import sys from collections.abc import Callable, Sequence from dataclasses import dataclass, replace @@ -17,6 +18,7 @@ from flashdreams.core.distributed import init as distributed_init from flashdreams.runtime.demo.bootstrap import ( + cleanup_cuda_distributed, configure_logging, initialize_cuda_distributed, ) @@ -228,6 +230,7 @@ def run_application_webrtc( manager=manager, runtime=runtime, primary_error=exc, + world_rank=context.world_rank, ) raise @@ -289,18 +292,64 @@ def _cleanup_application_webrtc_startup_failure( manager: Any | None, runtime: Any, primary_error: BaseException, + world_rank: int, ) -> None: - try: - if manager is None: - runtime.close() - return - asyncio.run(manager.shutdown()) - except BaseException as cleanup_error: - add_note = getattr(primary_error, "add_note", None) + errors: list[BaseException] = [] + if manager is None: + _record_application_webrtc_startup_cleanup_error(errors, runtime.close) + else: + _record_application_webrtc_startup_cleanup_error( + errors, + _send_application_webrtc_exit_signal, + manager, + ) + _record_application_webrtc_startup_cleanup_error( + errors, + _shutdown_application_webrtc_manager, + manager, + ) + _record_application_webrtc_startup_cleanup_error( + errors, + cleanup_cuda_distributed, + world_rank=world_rank, + synchronize_distributed=False, + torch_module=torch, + dist_module=dist, + ) + add_note = getattr(primary_error, "add_note", None) + for cleanup_error in errors: if callable(add_note): add_note(f"Additional WebRTC startup cleanup error: {cleanup_error!r}") +def _record_application_webrtc_startup_cleanup_error( + errors: list[BaseException], + cleanup: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, +) -> None: + try: + cleanup(*args, **kwargs) + except BaseException as cleanup_error: + errors.append(cleanup_error) + + +def _send_application_webrtc_exit_signal(manager: Any) -> None: + send_exit_signal = getattr(manager, "send_exit_signal", None) + if callable(send_exit_signal): + send_exit_signal() + + +def _shutdown_application_webrtc_manager(manager: Any) -> None: + shutdown = getattr(manager, "shutdown", None) + if not callable(shutdown): + return + result = shutdown() + if inspect.isawaitable(result): + asyncio.run(result) + + def _create_application_webrtc_resources( *, adapter: DemoAdapter, diff --git a/flashdreams/flashdreams/runtime/demo/bootstrap.py b/flashdreams/flashdreams/runtime/demo/bootstrap.py index 6b1e68f5d..a2b707dbc 100644 --- a/flashdreams/flashdreams/runtime/demo/bootstrap.py +++ b/flashdreams/flashdreams/runtime/demo/bootstrap.py @@ -5,6 +5,7 @@ from __future__ import annotations +import gc import logging import os from collections.abc import Callable @@ -91,8 +92,44 @@ def initialize_cuda_distributed( ) +def cleanup_cuda_distributed( + *, + world_rank: int | None = None, + synchronize_distributed: bool = True, + torch_module: Any = torch, + dist_module: Any = dist, +) -> None: + """Release process-level CUDA and distributed state owned by demo serving.""" + gc.collect() + cuda = torch_module.cuda + if cuda.is_available(): + cuda.empty_cache() + cuda.synchronize() + + if not _dist_is_initialized(dist_module): + return + if synchronize_distributed: + dist_module.barrier() + if world_rank is None: + logging.getLogger(__name__).info("Destroying process group.") + else: + logging.getLogger(__name__).info( + "[Rank %s] Destroying process group.", + world_rank, + ) + dist_module.destroy_process_group() + + +def _dist_is_initialized(dist_module: Any) -> bool: + is_available = getattr(dist_module, "is_available", None) + if callable(is_available) and not is_available(): + return False + return bool(dist_module.is_initialized()) + + __all__ = [ "DistributedDemoContext", + "cleanup_cuda_distributed", "configure_logging", "initialize_cuda_distributed", ] diff --git a/flashdreams/flashdreams/serving/webrtc/bootstrap.py b/flashdreams/flashdreams/serving/webrtc/bootstrap.py index 353e1a88c..ae0721824 100644 --- a/flashdreams/flashdreams/serving/webrtc/bootstrap.py +++ b/flashdreams/flashdreams/serving/webrtc/bootstrap.py @@ -5,7 +5,10 @@ from __future__ import annotations -import gc +import asyncio +import inspect +from collections.abc import Callable +from typing import Any import torch import torch.distributed as dist @@ -16,6 +19,7 @@ DistributedDemoContext as WebRTCDistributedContext, ) from flashdreams.runtime.demo.bootstrap import ( + cleanup_cuda_distributed, configure_logging, initialize_cuda_distributed, ) @@ -31,25 +35,97 @@ def run_webrtc_server( port: int, ) -> None: """Serve on rank 0, idle on worker ranks, then tear the runtime down.""" - if world_rank == 0: - if app is None: - raise ValueError("Rank 0 requires an aiohttp app to serve.") - try: - web.run_app(app, host=host, port=port) - finally: - session_manager.send_exit_signal() - else: - try: - session_manager.wait_for_termination() - except KeyboardInterrupt: - logger.warning("Worker rank interrupted, shutting down.") - - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - - if dist.is_initialized(): - dist.barrier() - logger.info("[Rank {}] Destroying process group", world_rank) - dist.destroy_process_group() + primary_error: BaseException | None = None + completed = False + try: + if world_rank == 0: + if app is None: + raise ValueError("Rank 0 requires an aiohttp app to serve.") + try: + web.run_app(app, host=host, port=port) + finally: + session_manager.send_exit_signal() + else: + try: + session_manager.wait_for_termination() + except KeyboardInterrupt: + logger.warning("Worker rank interrupted, shutting down.") + completed = True + except BaseException as exc: + primary_error = exc + raise + finally: + _cleanup_webrtc_process( + session_manager=session_manager, + world_rank=world_rank, + synchronize_distributed=completed, + primary_error=primary_error, + ) + + +def _cleanup_webrtc_process( + *, + session_manager: WebRTCServerLifecycle, + world_rank: int, + synchronize_distributed: bool, + primary_error: BaseException | None, +) -> None: + errors: list[BaseException] = [] + if primary_error is not None: + _record_webrtc_process_cleanup_error( + errors, + _shutdown_webrtc_session_manager, + session_manager, + ) + _record_webrtc_process_cleanup_error( + errors, + cleanup_cuda_distributed, + world_rank=world_rank, + synchronize_distributed=synchronize_distributed, + torch_module=torch, + dist_module=dist, + ) + if primary_error is not None: + _record_webrtc_process_cleanup_notes(primary_error, errors) + return + _raise_first_webrtc_process_cleanup_error(errors) + + +def _shutdown_webrtc_session_manager(session_manager: WebRTCServerLifecycle) -> None: + shutdown = getattr(session_manager, "shutdown", None) + if not callable(shutdown): + return + result = shutdown() + if inspect.isawaitable(result): + asyncio.run(result) + + +def _record_webrtc_process_cleanup_error( + errors: list[BaseException], + cleanup: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, +) -> None: + try: + cleanup(*args, **kwargs) + except BaseException as cleanup_error: + errors.append(cleanup_error) + + +def _raise_first_webrtc_process_cleanup_error(errors: list[BaseException]) -> None: + if not errors: + return + first = errors[0] + _record_webrtc_process_cleanup_notes(first, errors[1:]) + raise first + + +def _record_webrtc_process_cleanup_notes( + primary_error: BaseException, + errors: list[BaseException], +) -> None: + add_note = getattr(primary_error, "add_note", None) + for cleanup_error in errors: + if callable(add_note): + add_note(f"Additional WebRTC process cleanup error: {cleanup_error!r}") diff --git a/flashdreams/tests/test_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py index e45768a62..6ba089d79 100644 --- a/flashdreams/tests/test_demo_application_api.py +++ b/flashdreams/tests/test_demo_application_api.py @@ -957,6 +957,12 @@ def test_run_application_webrtc_closes_runtime_when_server_startup_fails( "initialize_cuda_distributed", lambda **_: SimpleNamespace(device="cuda:0", world_rank=0), ) + cleanup_calls: list[dict[str, object]] = [] + monkeypatch.setattr( + demo_app, + "cleanup_cuda_distributed", + lambda **kwargs: cleanup_calls.append(kwargs), + ) import flashdreams.serving.webrtc.demo as webrtc_demo @@ -970,6 +976,12 @@ def fail_to_serve(**_: object) -> object: assert len(adapter.runtimes) == 1 assert adapter.runtimes[0].closed + assert len(cleanup_calls) == 1 + cleanup_call = cleanup_calls[0] + assert cleanup_call["world_rank"] == 0 + assert cleanup_call["synchronize_distributed"] is False + assert cleanup_call["torch_module"] is demo_app.torch + assert cleanup_call["dist_module"] is demo_app.dist def test_demo_application_can_be_built_from_callbacks() -> None: diff --git a/flashdreams/tests/test_webrtc_bootstrap.py b/flashdreams/tests/test_webrtc_bootstrap.py index c83d4aac1..f66e29460 100644 --- a/flashdreams/tests/test_webrtc_bootstrap.py +++ b/flashdreams/tests/test_webrtc_bootstrap.py @@ -6,6 +6,7 @@ import pytest import torch +from flashdreams.runtime.demo import bootstrap as runtime_bootstrap from flashdreams.serving.webrtc import bootstrap pytestmark = pytest.mark.ci_cpu @@ -118,3 +119,114 @@ def test_initialize_cuda_distributed_rejects_cpu_default_device( with pytest.raises(RuntimeError, match="CUDA device is required"): bootstrap.initialize_cuda_distributed(default_device="cpu") + + +def test_cleanup_cuda_distributed_destroys_group_without_barrier() -> None: + fake_torch = _FakeTorch() + fake_dist = _FakeDist() + + runtime_bootstrap.cleanup_cuda_distributed( + world_rank=0, + synchronize_distributed=False, + torch_module=fake_torch, + dist_module=fake_dist, + ) + + assert fake_torch.cuda.empty_cache_calls == 1 + assert fake_torch.cuda.synchronize_calls == 1 + assert fake_dist.barrier_calls == 0 + assert fake_dist.destroy_process_group_calls == 1 + + +def test_run_webrtc_server_cleans_process_state_when_startup_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[str] = [] + manager = _FakeServerLifecycle(events) + cleanup_calls: list[dict[str, object]] = [] + + def fail_to_run_app(*_: object, **__: object) -> None: + raise RuntimeError("server startup failed") + + def record_cleanup(**kwargs: object) -> None: + events.append("process-cleanup") + cleanup_calls.append(kwargs) + + monkeypatch.setattr(bootstrap.web, "run_app", fail_to_run_app) + monkeypatch.setattr(bootstrap, "cleanup_cuda_distributed", record_cleanup) + + with pytest.raises(RuntimeError, match="server startup failed"): + bootstrap.run_webrtc_server( + world_rank=0, + session_manager=manager, + app=bootstrap.web.Application(), + host="127.0.0.1", + port=8080, + ) + + assert manager.send_exit_signal_calls == 1 + assert events == ["send-exit", "shutdown", "process-cleanup"] + assert len(cleanup_calls) == 1 + cleanup_call = cleanup_calls[0] + assert cleanup_call["world_rank"] == 0 + assert cleanup_call["synchronize_distributed"] is False + assert cleanup_call["torch_module"] is bootstrap.torch + assert cleanup_call["dist_module"] is bootstrap.dist + + +class _FakeCuda: + def __init__(self) -> None: + self.empty_cache_calls = 0 + self.synchronize_calls = 0 + + def is_available(self) -> bool: + return True + + def empty_cache(self) -> None: + self.empty_cache_calls += 1 + + def synchronize(self) -> None: + self.synchronize_calls += 1 + + +class _FakeTorch: + def __init__(self) -> None: + self.cuda = _FakeCuda() + + +class _FakeDist: + def __init__(self) -> None: + self.barrier_calls = 0 + self.destroy_process_group_calls = 0 + + def is_available(self) -> bool: + return True + + def is_initialized(self) -> bool: + return True + + def barrier(self) -> None: + self.barrier_calls += 1 + + def destroy_process_group(self) -> None: + self.destroy_process_group_calls += 1 + + +class _FakeServerLifecycle: + def __init__(self, events: list[str] | None = None) -> None: + self.events = events + self.send_exit_signal_calls = 0 + self.shutdown_calls = 0 + + def send_exit_signal(self) -> None: + self.send_exit_signal_calls += 1 + if self.events is not None: + self.events.append("send-exit") + + def wait_for_termination(self) -> None: + raise AssertionError("rank 0 should not wait for termination") + + async def shutdown(self) -> None: + self.shutdown_calls += 1 + if self.events is not None: + self.events.append("shutdown")