diff --git a/apps/t2v/README.md b/apps/t2v/README.md new file mode 100644 index 000000000..625b1084e --- /dev/null +++ b/apps/t2v/README.md @@ -0,0 +1,106 @@ +# 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. + +## 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 +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/apps/t2v/__init__.py b/apps/t2v/__init__.py new file mode 100644 index 000000000..12ecf05d3 --- /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, + model_config_from_runner, + 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", + "model_config_from_runner", + "t2v_scenario_mapping", +] diff --git a/apps/t2v_demo/pyproject.toml b/apps/t2v/pyproject.toml similarity index 61% rename from apps/t2v_demo/pyproject.toml rename to apps/t2v/pyproject.toml index e6b9f10e9..4521941aa 100644 --- a/apps/t2v_demo/pyproject.toml +++ b/apps/t2v/pyproject.toml @@ -6,21 +6,18 @@ requires = ["setuptools>=69", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "flashdreams-t2v-demo" +name = "flashdreams-t2v" version = "0.1.0" -description = "FlashDreams text-to-video runtime demo launcher" +description = "Shared model-neutral text-to-video demo application shell" requires-python = ">=3.10" dependencies = ["flashdreams[serving]"] -[project.entry-points."flashdreams.runner_configs"] -t2v = "t2v_demo.runner:RUNNER_T2V" - [tool.uv.sources] flashdreams = { workspace = true } [tool.setuptools] -packages = ["t2v_demo"] -package-dir = { t2v_demo = "." } +packages = ["t2v"] +package-dir = { t2v = "." } [tool.setuptools.package-data] -t2v_demo = ["web/*.js", "web/*.css"] +t2v = ["web/*.js", "web/*.css"] diff --git a/apps/t2v/t2v.py b/apps/t2v/t2v.py new file mode 100644 index 000000000..4b0ed5402 --- /dev/null +++ b/apps/t2v/t2v.py @@ -0,0 +1,706 @@ +# 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 + +import io +import json +import zipfile +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from importlib.resources import files +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import torch + +from flashdreams.demo import ( + Application, + DemoAdapterApplication, +) +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, + StepRequirements, +) +from flashdreams.runtime._utils import freeze_mapping +from flashdreams.runtime.demo import ( + DemoSpec, + NullOutputSpec, + OutputSpec, + PreparedScenario, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.runtime.demo.outputs import SessionInfo +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" + +_DOWNLOAD_ARTIFACT_DIR = Path("outputs/t2v-webrtc") + + +@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.") + # 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: + """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) + ) + + 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: + """Prompt and output geometry for a finite text-to-video rollout.""" + + prompt: str + total_blocks: int + pixel_height: int + 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.""" + + 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, write_download_artifact: bool = False + ) -> None: + self.model = model + self.write_download_artifact = write_download_artifact + + @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}, 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: + 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=scenario.to_mapping()) + ) + + def create_runtime(self, config: InferenceConfig) -> "T2VRuntime": + self.validate_config(config) + return T2VRuntime( + config=config, + model=self.model, + write_download_artifact=self.write_download_artifact, + ) + + 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) + + 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.""" + + 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: StepRequirements, 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, + 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 + + 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=T2VScenario.from_mapping(inputs.global_conditioning), + runtime=self, + artifact_dir=( + _DOWNLOAD_ARTIFACT_DIR if self._write_download_artifact else None + ), + ) + + 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, + artifact_dir: Path | None = None, + ) -> None: + self.pipeline = pipeline + self.scenario = scenario + self._runtime = runtime + # 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" + ) + 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( + "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=self.scenario.to_mapping()) + + def next_step_requirements(self) -> StepRequirements | None: + if self._closed or self._step_index >= self.scenario.total_blocks: + return None + 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 + 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}, + ) + if self._artifact_output is not None: + self._artifact_output.write(result) + return result + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + raise RuntimeError( + "T2V sessions are finite; create a new session instead of reset()." + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + 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) + + +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, write_download_artifact=output.mode == "webrtc" + ), + 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, + 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, + }, + ), + metadata={ + "output_layout": "tchw", + "webrtc_keep_connection_after_completed": True, + "webrtc_preload_name": "FlashDreams T2V", + "webrtc_supported_control_keys": ("g",), + }, + ) + + +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 T2VScenario.from_mapping( + defaults.scenario_overrides(), defaults=model.default_scenario() + ).to_mapping() + + +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) + + +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", + "FIELD_PIXEL_WIDTH", + "FIELD_PROMPT", + "FIELD_TOTAL_BLOCKS", + "T2VDemoAdapter", + "T2VInputProvider", + "T2VModelConfig", + "T2VRunDefaults", + "T2VRuntime", + "T2VScenario", + "T2VSession", + "create_t2v_application", + "create_t2v_spec", + "model_config_from_runner", + "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..666968fbd --- /dev/null +++ b/apps/t2v/tests/test_t2v_shell.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import inspect +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from t2v.t2v import ( + T2VDemoAdapter, + T2VModelConfig, + T2VRunDefaults, + create_t2v_application, + create_t2v_spec, + model_config_from_runner, + t2v_scenario_mapping, +) + +from flashdreams.demo import DemoAdapterApplication +from flashdreams.runtime.demo import ( + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) + +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_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_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 + + source = inspect.getsource(t2v_shell) + + 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/web/adapter.css b/apps/t2v/web/adapter.css similarity index 100% rename from apps/t2v_demo/web/adapter.css rename to apps/t2v/web/adapter.css diff --git a/apps/t2v_demo/web/adapter.js b/apps/t2v/web/adapter.js similarity index 100% rename from apps/t2v_demo/web/adapter.js rename to apps/t2v/web/adapter.js 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 bee09cb67..000000000 --- a/apps/t2v_demo/app.py +++ /dev/null @@ -1,345 +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 flashdreams.runtime import InferenceConfig -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.runtime.demo.replay import run_replay_demo -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 .runtime import ( - FIELD_FPS, - FIELD_PIXEL_HEIGHT, - FIELD_PIXEL_WIDTH, - FIELD_PROMPT, - FIELD_TOTAL_BLOCKS, - T2VDemoAdapter, - make_adapter, -) - -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 - - -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 {} - adapter = make_adapter(config.backend) - scenario = _scenario(config, scenario_overrides) - 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_replay_demo( - spec=_spec( - config, - adapter=adapter, - scenario=scenario, - input_mode="replay", - output=output, - ), - adapter=adapter, - ) - 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) - 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( - config, - adapter=adapter, - scenario=scenario, - input_mode="webrtc", - output=output, - device=str(context.device), - ) - 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 _scenario( - config: "T2VDemoRunnerConfig", overrides: dict[str, object] -) -> dict[str, object]: - runner = resolve_backend(config.backend).resolve_runner(config.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) - ) - - 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), - } - - -def _spec( - config: "T2VDemoRunnerConfig", - *, - 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, - 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, - runtime_options={"backend": adapter.backend.key}, - ), - ) - - -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 _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__ = ["T2VWebRTCSessionManager", "launch_t2v"] 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/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/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/runner.py b/apps/t2v_demo/runner.py deleted file mode 100644 index 8fe620568..000000000 --- a/apps/t2v_demo/runner.py +++ /dev/null @@ -1,72 +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 .backends import backend_choices, resolve_backend - - -@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=lambda: resolve_backend("causal-forcing") - .resolve_runner() - .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 873496c17..000000000 --- a/apps/t2v_demo/runtime.py +++ /dev/null @@ -1,297 +0,0 @@ -# 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.""" - -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 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 make_adapter(backend: str) -> T2VDemoAdapter: - """Build an adapter from a CLI/UI backend key.""" - return T2VDemoAdapter(backend=resolve_backend(backend)) diff --git a/apps/t2v_demo/tests/test_runner.py b/apps/t2v_demo/tests/test_runner.py deleted file mode 100644 index 3bc17756e..000000000 --- a/apps/t2v_demo/tests/test_runner.py +++ /dev/null @@ -1,47 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import pytest -from t2v_demo import app -from t2v_demo.runner import RUNNER_T2V, T2VDemoRunnerConfig - -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_runner_mp4_launch_uses_demo_entrypoint( - monkeypatch: pytest.MonkeyPatch, -) -> None: - captured = [] - - def fake_replay_demo(*, spec: object, adapter: object) -> object: - captured.append((spec, adapter)) - return type("Result", (), {"status": "completed"})() - - monkeypatch.setattr(app, "run_replay_demo", fake_replay_demo) - 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}, - ) - - 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 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 new file mode 100644 index 000000000..a2b91df0f --- /dev/null +++ b/flashdreams/flashdreams/demo/__init__.py @@ -0,0 +1,90 @@ +# 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.app import ( + DemoApplication, + create_demo_application, + run_application_replay, + run_application_webrtc, + run_replay_application, +) +from flashdreams.demo.application import ( + Application, + ApplicationSession, + DemoAdapterApplication, + FrameOutputSink, + IApplication, + IApplicationSession, + InferenceSessionApplicationAdapter, + IOHandler, + 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, + NativeWindowIOHandler, + ReplayIOHandler, + WebRTCIOHandlerServer, + create_native_window_io_handler, + create_replay_io_handler, + 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", + "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_demo_application", + "create_replay_io_handler", + "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 new file mode 100644 index 000000000..a8d720a9e --- /dev/null +++ b/flashdreams/flashdreams/demo/app.py @@ -0,0 +1,408 @@ +# 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 asyncio +import inspect +import sys +from collections.abc import Callable, Sequence +from dataclasses import dataclass, replace +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 ( + cleanup_cuda_distributed, + 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, + 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.""" + + 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() + 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) + + 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) + + 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) + + 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.""" + 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.""" + 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.") + + 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 run_application_replay( + app=DemoAdapterApplication(adapter=adapter, spec=spec) + ) + + +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): + output_sink = build_output_sink(app.spec.output) + return Runner( + io_handler=create_replay_io_handler(output_sink=output_sink), + app=app, + launch_args=tuple(launch_args), + ).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: 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 + + 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, + world_rank=context.world_rank, + ) + raise + + +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 _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 _cleanup_application_webrtc_startup_failure( + *, + manager: Any | None, + runtime: Any, + primary_error: BaseException, + world_rank: int, +) -> 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, + 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 + 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", + "create_demo_application", + "run_application_webrtc", + "run_application_replay", + "run_replay_application", +] diff --git a/flashdreams/flashdreams/demo/application.py b/flashdreams/flashdreams/demo/application.py new file mode 100644 index 000000000..33b23fae4 --- /dev/null +++ b/flashdreams/flashdreams/demo/application.py @@ -0,0 +1,293 @@ +# 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.config import InferenceConfig +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.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, +) + +from .inputs import InputName + + +@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. + """ + ... + + def close(self) -> None: + """Release application-level resources after all sessions have closed.""" + ... + + +@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: InputName | 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) + _runtimes: list[InferenceRuntime] = field( + default_factory=list, + init=False, + repr=False, + ) + + def init(self, launch_args: Sequence[str]) -> None: + 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) + + 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._runtimes.append(runtime) + return InferenceSessionApplicationAdapter( + 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: + 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) +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/flashdreams/demo/inputs.py b/flashdreams/flashdreams/demo/inputs.py new file mode 100644 index 000000000..c32d50302 --- /dev/null +++ b/flashdreams/flashdreams/demo/inputs.py @@ -0,0 +1,265 @@ +# 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, + WSAD_SUPPORTED_KEYS, + KeyboardState, + 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 new file mode 100644 index 000000000..c07352c2d --- /dev/null +++ b/flashdreams/flashdreams/demo/io.py @@ -0,0 +1,482 @@ +# 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 ( + CompositeOutputSinkError, + 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 +from .inputs import ( + InputName, + InputStateDecoderRegistry, + create_default_input_state_decoder_registry, +) + +RunSessionCallback = Callable[[IOHandler], RunResult] +ServeCallback = Callable[[], object] + + +@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_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, + 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: + 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.""" + 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 + 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) + 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 + 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) + 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 + 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 + return decision + + def should_exit(self) -> bool: + return self._should_exit or self._closed + + def close(self) -> Sequence[OutputArtifact]: + self._closed = True + self._should_exit = True + 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) +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 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.""" + + 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 _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, + 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 + + +def _coerce_run_result(value: object) -> RunResult: + if isinstance(value, RunResult): + return value + return RunResult(status="completed") + + +__all__ = [ + "CallbackIOHandlerServer", + "IOHandlerBatchInputSource", + "IOHandlerOutputSink", + "IOHandlerRunMode", + "IOHandlerServer", + "NativeWindowIOHandler", + "ReplayIOHandler", + "RunSessionCallback", + "ServeCallback", + "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 new file mode 100644 index 000000000..172875f05 --- /dev/null +++ b/flashdreams/flashdreams/demo/runner.py @@ -0,0 +1,638 @@ +# 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 Callable, Coroutine, 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 ( + 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 +from flashdreams.runtime.demo.run_modes import ( + InMemorySessionMetricsRecorder, + RunContext, + RunMode, + RunResult, + 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, +) +from flashdreams.runtime.interfaces import InferenceRuntime, InferenceSession +from flashdreams.runtime.mapping import InputMapping +from flashdreams.runtime.types import StepRequirements + +from .application import ( + Application, + ApplicationSession, + DemoAdapterApplication, + IOHandler, +) +from .io import IOHandlerRunMode, ReplayIOHandler + + +@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() + spec = self._create_spec(run_mode) + context = self._create_context(host) + 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) + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + result = await run_demo_session_async( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=run_mode, + pipeline=self.pipeline or StepPipeline(), + ) + return result + except BaseException as exc: + primary_error = exc + raise + finally: + await _await_runner_cleanup( + _close_runner_resources_async( + context=context, + host=host, + 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, + ), + 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: + selected_run_mode = run_mode or self._selected_run_mode() + host, owns_host = self._selected_host() + spec = self._create_spec(selected_run_mode) + context = self._create_context(host) + 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) + adapter = _RunnerDemoAdapter(app=self.app, spec=spec, scenario=scenario) + result = run_demo_session( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=selected_run_mode, + pipeline=self.pipeline or StepPipeline(), + ) + return result + except BaseException as exc: + primary_error = exc + raise + finally: + _close_runner_resources( + context=context, + host=host, + 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, + ) + + def _selected_run_mode(self) -> RunMode: + if self.run_mode is not None: + return self.run_mode + 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: + 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: + 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, + input_mode=run_mode.name, + output=NullOutputSpec(), + 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: + 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, " + f"got {type(session).__name__}." + ) + session.init() + return cast(InferenceSession, session) + + def close(self) -> None: + return None + + +@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: + 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}.") + + def create_runtime(self, config: InferenceConfig) -> InferenceRuntime: + self.validate_config(config) + 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: + if spec != self.spec: + raise ValueError("Runner received an unexpected DemoSpec.") + return self.scenario + + def create_model_input_provider( + self, + spec: DemoSpec, + scenario: Any, + ) -> 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() + + +@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 _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) + + +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()) + + +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, + host: RuntimeHost, + 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 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: + _record_cleanup_notes(primary_error, errors) + return + 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 _close_runner_resources_async( + *, + context: RunContext, + host: RuntimeHost, + 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] = [] + try: + await context.close_async() + except Exception as exc: + errors.append(exc) + 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: + _record_cleanup_notes(primary_error, errors) + return + 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], + *, + preserve_primary: bool = False, + preserved_error: BaseException | None = 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: + 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 + + +def _record_cleanup_error( + errors: list[Exception], + cleanup: Any, + /, + *args: Any, +) -> None: + try: + cleanup(*args) + except Exception as exc: + errors.append(exc) + + +def _close_application( + *, + errors: list[Exception], + host: RuntimeHost, + cleanup: _ApplicationCleanup, +) -> None: + if cleanup.closed: + return + invoked = False + + def close_app() -> None: + nonlocal invoked + invoked = True + 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 + errors.append(exc) + + +async def _close_application_async( + *, + errors: list[Exception], + host: RuntimeHost, + cleanup: _ApplicationCleanup, +) -> None: + if cleanup.closed: + return + invoked = False + + def close_app() -> None: + nonlocal invoked + invoked = True + 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 + errors.append(exc) + + +def _raise_first_cleanup_error(errors: Sequence[Exception]) -> None: + if not errors: + return + first = errors[0] + _record_cleanup_notes(first, errors[1:]) + raise first + + +def _record_cleanup_notes( + primary: BaseException | None, + errors: Sequence[BaseException], +) -> 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}") + + +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 _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: + # 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 " + "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 + + +def _application_adapter(app: Application) -> DemoAdapter | None: + if isinstance(app, DemoAdapterApplication): + return app.adapter + return None + + +__all__ = ["Runner"] 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/flashdreams/runtime/demo/__init__.py b/flashdreams/flashdreams/runtime/demo/__init__.py index bf64905b2..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, @@ -57,6 +60,8 @@ RunSummary, SessionDriver, SessionEdges, + SessionExitSource, + SessionExitState, SingleSessionAdmissionPolicy, WebRTCErrorPolicy, build_model_warmup_plan, @@ -131,6 +136,8 @@ "BenchmarkStatsOutputSink", "BenchmarkRunMode", "BenchmarkErrorPolicy", + "ComparisonOutputMismatchError", + "ComparisonOutputSink", "InMemorySessionMetricsRecorder", "InputSource", "CatchUpDecision", @@ -138,6 +145,7 @@ "CompositeOutputSink", "CompositeOutputSinkError", "DeterministicClock", + "FileOutputSink", "MetricsSnapshot", "ModelWarmupAdapter", "ModelWarmupPlan", @@ -175,6 +183,8 @@ "RuntimeHost", "SessionEdges", "SessionDriver", + "SessionExitState", + "SessionExitSource", "SessionInfo", "SignalActivationPolicy", "SingleSessionAdmissionPolicy", 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/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/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/host.py b/flashdreams/flashdreams/runtime/demo/host.py index 6a6e93ce5..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 @@ -81,6 +84,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.""" @@ -135,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() @@ -154,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() @@ -170,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/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/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/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 c366debcd..fd794b917 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,21 @@ 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, + run_application_webrtc, +) from flashdreams.infra.runner import RunnerConfig +from flashdreams.plugins import discover_applications +from flashdreams.runtime.demo import ( + DemoAdapter, + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + WebRTCOutputSpec, +) from flashdreams.serving.launch import ( LaunchMode, LaunchOptions, @@ -199,6 +213,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 +350,481 @@ 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", "webrtc"}: + raise ValueError( + 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_launch( + 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 + 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( + 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, 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] = [] + 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." + ) + 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_launch( + *, + 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, + ) + adapter = _application_adapter_for_output(application.adapter, output) + return DemoAdapterApplication( + adapter=adapter, + spec=dataclasses.replace( + application.spec, + 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], +) -> 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 | 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 == "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"}, + ) + 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 _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) + 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 _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], + *, + 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", "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 mode settings.") + + def _prepare_cli_args( args: list[str], ) -> tuple[ 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/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 149f83d4f..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.""" @@ -728,6 +753,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(): @@ -1628,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_demo_application_api.py b/flashdreams/tests/test_demo_application_api.py new file mode 100644 index 000000000..6ba089d79 --- /dev/null +++ b/flashdreams/tests/test_demo_application_api.py @@ -0,0 +1,1854 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import asyncio +import threading +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, + ApplicationSession, + BenchmarkStatsOutputSink, + CallbackIOHandlerServer, + ComparisonOutputSink, + DemoAdapterApplication, + DemoApplication, + FileOutputSink, + FrameOutputSink, + InferenceSessionApplicationAdapter, + InputName, + IOHandler, + IOHandlerServer, + KeyboardInputState, + ReplayIOHandler, + Runner, + RuntimeOutputSinkFrameAdapter, + WebRTCIOHandlerServer, + create_demo_application, + create_native_window_io_handler, + create_replay_io_handler, + create_webrtc_io_handler, + input_state_from_window, + run_application_replay, + run_application_webrtc, +) +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InferenceSession, + InputField, + InputMapping, + StepRequest, + StepResult, + TimeWindow, + UserInputCapability, + UserInputEvent, + UserInputs, + UserInputSchema, +) +from flashdreams.runtime.demo import ( + BatchSessionDriver, + DemoSpec, + InMemorySessionMetricsRecorder, + NullOutputSink, + NullOutputSpec, + OutputDecision, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + RunContext, + RunModeCapabilities, + RunResult, + RuntimeHost, + SessionEdges, + SessionInfo, + SingleSessionAdmissionPolicy, + StepPipeline, + UserInputWindow, + WebRTCOutputSpec, +) +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 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" + 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__ == ( + "create_native_window_io_handler" + ) + assert create_webrtc_io_handler.__name__ == "create_webrtc_io_handler" + + +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: + adapter = _FakeDemoAdapter() + demo = DemoAdapterApplication( + adapter=adapter, + 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() + 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: + 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_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_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()) + 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 + + +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 app.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 + + +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_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_does_not_init_public_app_when_host_starts_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.init_thread_id is None + assert not app.closed + + +def test_runner_raises_when_external_host_disappears_before_cleanup() -> None: + app = _RunnerFakeApplication(total_steps=1) + host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) + + try: + 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 host.is_closed + assert not 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 + + +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) + + 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_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) + + 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) + 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] + + +@pytest.mark.asyncio +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() + + 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.init_thread_id is None + assert not app.closed + + +@pytest.mark.asyncio +async def test_runner_run_async_raises_when_external_host_disappears_before_cleanup() -> ( + None +): + app = _RunnerFakeApplication(total_steps=1) + host = _ExternallyClosedWithoutHooksRuntimeHost(_ExternalApplicationRuntime(app)) + + try: + 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 host.is_closed + 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) + + 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 + + +@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) + + 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_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) + + 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 + + +@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 + + +@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) + 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_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_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) + + 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_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)) + + 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) + + +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_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_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), + ) + 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 + + 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 + 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: + 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() + + app.main(["webrtc"]) + + assert app.server_called + 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}, + ) + + +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: + return _parse_replay_command(argv) + + def replay_spec(self, args: argparse.Namespace) -> DemoSpec: + return _fake_replay_spec(args) + + 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"),) + ) + + 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 + + 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=self.step_index, + inference_input_schema=self.inference_input_schema, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + self.inference_input_schema.require_step(inputs) + 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 + 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 _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( + global_conditioning_fields=(InputField(name="prompt"),), + step_fields=(InputField(name="chunk_index"),), + ) + canonical_input_schema = CanonicalInputSchema() + + 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",) + + 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 = _FakeRuntime() + 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=self.user_inputs, + source_schema=self.source_schema, + ) + + 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: + 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 () + + +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 _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" + + def __init__( + self, + *, + 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 + self.closed = False + + 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") + + def create_session(self) -> "_RunnerFakeSession": + self.session = _RunnerFakeSession( + total_steps=self.total_steps, + fail_step=self.fail_step, + ) + return self.session + + def close(self) -> None: + self.closed = True + + +class _RunnerPrimaryBaseException(BaseException): + pass + + +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 + self.closed = False + + def start_session(self, inputs: InferenceInput) -> InferenceSession: + del inputs + session = self._app.create_session() + session.init() + return _ExternalInferenceSession(session) + + 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") + + +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, 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, ...] = () + 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) + 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, + 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 () + + +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, + ) + + +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") + + +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 + name: str = "async-public-runner" + capabilities: RunModeCapabilities = field( + default_factory=lambda: RunModeCapabilities(supports_artifacts=True) + ) + driver: Any = 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) -> Any: + 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 _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 _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 _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 + 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() 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()) 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) diff --git a/flashdreams/tests/test_launch_manifest.py b/flashdreams/tests/test_launch_manifest.py index 7adbdd152..bbcf6537d 100644 --- a/flashdreams/tests/test_launch_manifest.py +++ b/flashdreams/tests/test_launch_manifest.py @@ -9,7 +9,25 @@ import pytest +from flashdreams.demo import DemoAdapterApplication from flashdreams.infra.runner import RunnerConfig +from flashdreams.runtime import ( + CanonicalInputSchema, + IdentityInputMapping, + InferenceConfig, + InferenceInput, + InferenceInputSchema, + InferenceRuntime, + InputMapping, +) +from flashdreams.runtime.demo import ( + DemoSpec, + Mp4OutputSpec, + NullOutputSpec, + PreparedScenario, + RunResult, + WebRTCOutputSpec, +) from flashdreams.scripts import cli from flashdreams.serving.launch import ResolvedLaunch, resolve_launch from flashdreams.serving.launch_manifest import load_launch_manifest @@ -29,6 +47,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 +89,110 @@ 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_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, @@ -441,3 +577,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", "webrtc") + + def supported_output_modes(self) -> tuple[str, ...]: + return ("mp4", "null", "webrtc") + + 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) -> InferenceRuntime: + del config + raise AssertionError("direct app CLI test should not instantiate the runtime") 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/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/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") 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: 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/causal_forcing/causal_forcing/t2v/__init__.py b/integrations/causal_forcing/causal_forcing/t2v/__init__.py new file mode 100644 index 000000000..0507d5c35 --- /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, 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 new file mode 100644 index 000000000..631e4cc0e --- /dev/null +++ b/integrations/causal_forcing/causal_forcing/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 Causal-Forcing model.""" + +from t2v import create_t2v_application, model_config_from_runner + +from causal_forcing.config import RUNNER_WAN21_T2V_1PT3B_CHUNKWISE +from flashdreams.demo import Application + +MODEL = model_config_from_runner( + model_id="causal-forcing-t2v", + runner=RUNNER_WAN21_T2V_1PT3B_CHUNKWISE, +) + + +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..644d23d78 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 RUNNER_CONFIGS, RUNNER_WAN21_T2V_1PT3B_CHUNKWISE +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 == 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 + + +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/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/cosmos_predict2/cosmos_predict2/t2v/__init__.py b/integrations/cosmos_predict2/cosmos_predict2/t2v/__init__.py new file mode 100644 index 000000000..68fcfb392 --- /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, 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 new file mode 100644 index 000000000..8e2716ff7 --- /dev/null +++ b/integrations/cosmos_predict2/cosmos_predict2/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 Cosmos Predict2 model.""" + +from t2v import create_t2v_application, model_config_from_runner + +from cosmos_predict2.config import RUNNER_COSMOS2_T2V_2B_720P +from flashdreams.demo import Application + +MODEL = model_config_from_runner( + model_id="cosmos-predict2-t2v", + runner=RUNNER_COSMOS2_T2V_2B_720P, +) + + +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..1a50b47ed --- /dev/null +++ b/integrations/cosmos_predict2/tests/test_t2v_app.py @@ -0,0 +1,41 @@ +# 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 RUNNER_COSMOS2_T2V_2B_720P +from cosmos_predict2.t2v.app import MODEL, create_app, createApp + +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 == 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 + + +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/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/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/lingbot/lingbot/demo/app.py b/integrations/lingbot/lingbot/demo/app.py index c47abf9b9..7841e773c 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 create_demo_application, 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, @@ -127,28 +127,19 @@ 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 prepare_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: +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, ) - - def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: from .webrtc import serve_lingbot_webrtc_demo - serve_lingbot_webrtc_demo( + return serve_lingbot_webrtc_demo( spec=_webrtc_spec( args, device=str(context.device), @@ -157,13 +148,7 @@ def serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: world_rank=context.world_rank, ) - -_APPLICATION = LingbotDemoApplication() - - -def main(argv: list[str] | None = None) -> None: - """Run the Lingbot demo application.""" - _APPLICATION.main(argv) + return CallbackIOHandlerServer(serve) def launch_from_runner( @@ -230,7 +215,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}.") @@ -386,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/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..534be6b33 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 create_demo_application, 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, @@ -125,33 +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 serve_webrtc(self, args: argparse.Namespace, *, context: Any) -> None: +def _webrtc_io_handler( + args: argparse.Namespace, + *, + context: Any, +) -> IOHandlerServer: + def serve() -> object: from .webrtc import serve_omnidreams_webrtc_demo - 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( @@ -226,7 +213,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}.") @@ -389,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) 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( 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/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/self_forcing/self_forcing/t2v/__init__.py b/integrations/self_forcing/self_forcing/t2v/__init__.py new file mode 100644 index 000000000..57b5c47fb --- /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, 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 new file mode 100644 index 000000000..bbb10e54c --- /dev/null +++ b/integrations/self_forcing/self_forcing/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 Self-Forcing model.""" + +from t2v import create_t2v_application, model_config_from_runner + +from flashdreams.demo import Application +from self_forcing.config import RUNNER_WAN21_T2V_1PT3B + +MODEL = model_config_from_runner( + model_id="self-forcing-t2v", + runner=RUNNER_WAN21_T2V_1PT3B, +) + + +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..b5c482d8c 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 RUNNER_CONFIGS, RUNNER_WAN21_T2V_1PT3B +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 == 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 + + +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/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/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/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 7f4aba302..32552fd6b 100644 --- a/uv.lock +++ b/uv.lock @@ -29,7 +29,7 @@ members = [ "flashdreams-omnidreams", "flashdreams-sana-wm", "flashdreams-self-forcing", - "flashdreams-t2v-demo", + "flashdreams-t2v", "flashdreams-wan21", "flashdreams-wan22", "ludus-renderer", @@ -1113,6 +1113,7 @@ version = "0.1.0" source = { editable = "integrations/causal_forcing" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, { name = "opencv-python-headless" }, ] @@ -1125,6 +1126,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" }, @@ -1137,6 +1139,7 @@ version = "0.1.0" source = { editable = "integrations/cosmos_predict2" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1148,6 +1151,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" }, ] @@ -1159,6 +1163,7 @@ version = "0.1.0" source = { editable = "integrations/fastvideo_causal_wan22" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1170,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" }, ] @@ -1402,6 +1408,7 @@ version = "0.1.0" source = { editable = "integrations/self_forcing" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, ] @@ -1413,15 +1420,16 @@ 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" }, ] provides-extras = ["dev"] [[package]] -name = "flashdreams-t2v-demo" +name = "flashdreams-t2v" version = "0.1.0" -source = { editable = "apps/t2v_demo" } +source = { editable = "apps/t2v" } dependencies = [ { name = "flashdreams", extra = ["serving"] }, ] @@ -1435,6 +1443,7 @@ version = "0.1.0" source = { editable = "integrations/wan21" } dependencies = [ { name = "flashdreams" }, + { name = "flashdreams-t2v" }, { name = "mediapy" }, { name = "opencv-python-headless" }, ] @@ -1447,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" },