diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md new file mode 100644 index 000000000..d621a0d33 --- /dev/null +++ b/PR_DESCRIPTION.md @@ -0,0 +1,133 @@ +## Summary + +Add a model-neutral `flashdreams-runner` shell and a `t2v-app` example +application for running FlashDreams pipelines through replay/MP4, WebRTC, or +headless I/O modes. + +The application owns inference. Its runtime owns model weights and one-time +initialization; each session owns its prompt, cache, step state, and generation +logic. The runner owns mode selection, process setup, lifecycle, the main loop, +and presentation. + +## High-level design + +```text +uv run flashdreams-runner t2v-app {mp4 | replay | webrtc | none} + | + v ++--------------------------+ create_runtime() +--------------------+ +| flashdreams_runner | ---------------------------> | t2v_app | +| | <--------------------------- | | +| select I/O mode | Runtime | Runtime | +| initialize Runtime | | model + pipeline | +| create Session | | | +| | input / output | Session | +| main loop: | ---------------------------> | prompt + cache | +| read input | <--------------------------- | generate/finalize | +| Session.step | StepResult +--------------------+ +| present output | ++------------+-------------+ + | + +-----+-----+----------------+ + | | | + v v v + MP4/replay WebRTC None +``` + +## Entrypoint examples + +Generate an MP4 with the packaged default preset: + +```bash +uv run flashdreams-runner t2v-app mp4 \ + --prompt "A waterfall" \ + --output o.mp4 +``` + +Use the explicit replay mode name and override its finite iteration count: + +```bash +uv run flashdreams-runner t2v-app replay \ + --steps 4 \ + --prompt "A waterfall" \ + --output o.mp4 +``` + +Serve the same application through WebRTC: + +```bash +uv run flashdreams-runner t2v-app webrtc \ + --prompt "A waterfall" +``` + +Run without presentation or artifacts: + +```bash +uv run flashdreams-runner t2v-app none \ + --steps 2 \ + --prompt "A waterfall" +``` + +Select a packaged preset explicitly: + +```bash +uv run flashdreams-runner t2v-app mp4 \ + --preset-id self-forcing-wan2.1-t2v-1.3b \ + --prompt "A neon-lit city at night" \ + --output outputs/city.mp4 +``` + +When `--preset-id` is omitted, `t2v-app` uses the catalog's +`default_preset_id`. The packaged default is +`causal-forcing-wan2.1-t2v-1.3b-chunkwise`. + +## Application ABI + +- Require an application module to expose only + `create_runtime(ApplicationArguments) -> Runtime`. +- Let the application factory extend the selected mode parser and resolve all + application-specific command-line configuration. +- Define `Runtime.initialize()`, `Runtime.create_session()`, and + `Runtime.destroy()` for one-time model and process state. +- Keep presentation fields in application-owned `AppConfig`, exposed through + `Runtime.config` for runner modes. +- Define `Session.step()` and `Session.destroy()` for per-user prompt, + cache, world state, and main-loop logic. +- Keep compatibility methods on the base runtime/session classes so shared + FlashDreams WebRTC code consumes application runtimes directly without a + runner-specific adapter. + +## Runner and modes + +- Add the root-level `flashdreams-runner` workspace package and console entrypoint. +- Select and construct runner-owned I/O modes independently of applications. +- Initialize the application runtime with the selected device and I/O handler. +- Own session creation, deterministic batch input, output delivery, and cleanup. +- Add finite replay/MP4, live WebRTC, and finite headless `none` modes behind an + extensible `IOHandler` contract. +- Keep `mp4` as a compatibility name for replay-to-file behavior. + +## T2V example application + +- Add `t2v-app` as an implementation of the application ABI. +- Resolve pipeline object graphs from packaged YAML presets without depending + on the legacy runner registry. +- Construct and retain the FlashDreams pipeline in `T2VRuntime`. +- Create the prompt-conditioned cache and run pipeline `generate`/`finalize` + inside `T2VSession.step()`. +- Keep prompts, dimensions, caches, and step indexes isolated per session. + +## Shared FlashDreams changes + +- Add reusable pipeline-preset parsing and provider loading under + `flashdreams.core.pipeline_presets`. +- Share the generator checkpoint prefix-remapping helper from FlashDreams core + across causal-forcing and self-forcing configurations. + +## Validation + +- Affected CPU tests cover the application ABI, runner lifecycle, mode + separation, WebRTC construction, T2V runtime/session ownership, and preset + resolution. +- Ruff, `ty`, Basedpyright, lockfile validation, and CLI help checks pass. +- GPU model generation was not run as part of this change. diff --git a/apps/t2v_app/ARCHITECTURE.md b/apps/t2v_app/ARCHITECTURE.md new file mode 100644 index 000000000..e1eaa8344 --- /dev/null +++ b/apps/t2v_app/ARCHITECTURE.md @@ -0,0 +1,133 @@ +# T2V Application Architecture + +The T2V application is an adapter between `flashdreams-runner` and a +FlashDreams streaming inference pipeline. The runner owns orchestration and +presentation; the application owns model configuration, model state, and +generation. + +## Minimal application contract + +The runner discovers the installed `t2v-app` distribution and imports its +top-level `t2v_app` module. That module exposes one public factory: + +```python +create_runtime(arguments: ApplicationArguments) -> Runtime +``` + +The returned runtime implements the contract in +[`flashdreams_runner/contracts.py`](../../flashdreams_runner/contracts.py): + +- `config` describes output identity, frame rate, layout, dimensions, and the + optional default step count. +- `initialize(device, io_handler)` constructs process-wide model state. +- `create_session(initial_input)` creates isolated generation state. +- `destroy()` releases process-wide resources. + +Each session implements: + +- `step_index`, the next autoregressive iteration. +- `step(inputs)`, which returns a FlashDreams `StepResult`. +- `destroy()`, which releases session state. + +`Runtime` and `Session` also adapt this runner-facing API to the shared +FlashDreams `InferenceRuntime` and `InferenceSession` protocols through +`start_session`, `next_step_request`, and `close`. + +## Components + +### Application factory + +[`t2v_app/application.py`](t2v_app/application.py) parses application +arguments, loads a pipeline preset, and creates an uninitialized `T2VRuntime`. +It does not construct model weights. + +### Runtime + +[`t2v_app/runtime.py`](t2v_app/runtime.py) owns the configured pipeline and +process-wide model weights. Initialization constructs the pipeline and moves it +to the selected device. The runtime creates one `T2VSession` for each isolated +generation. + +### Session + +[`t2v_app/session.py`](t2v_app/session.py) owns per-generation state: + +- prompt and video dimensions; +- autoregressive cache; +- current block index; +- optional WebRTC recording. + +Each generation step calls the pipeline's `generate` and `finalize` methods and +wraps the resulting video tensor in a `StepResult`. + +### WebRTC customization + +[`t2v_app/webrtc.py`](t2v_app/webrtc.py) is an optional adapter installed only +when the selected I/O handler is `WebRTCMode`. It supplies browser assets, +initial session input, prompt and duration updates, playback, and artifact +download routes. + +## Control flow + +### Startup + +```text +flashdreams-runner + -> import t2v_app + -> t2v_app.create_runtime(arguments) + -> T2VRuntime.initialize(device, io_handler) + -> io_handler.run(runtime, drive_session) +``` + +### Finite modes (`mp4`, `replay`, and `none`) + +```text +IO handler + -> drive_session(runtime, input_handler, output_handler) + -> runtime.create_session(initial_input) + -> session.step(step_input), repeated until input ends + -> output_handler.write(step_result) + -> session.destroy() +``` + +### WebRTC mode + +```text +WebRTCMode + -> T2VWebRTCCustomization.prepare_initial_input() + -> shared WebRTC session manager + -> runtime.start_session(initial_input) + -> session.next_step_request() + -> session.step(step_input), repeated until complete + -> session.close() +``` + +Browser prompt updates call `T2VRuntime.prepare_session_input()` to replace the +initial input used by the next generation. + +## Data boundary + +The primary values crossing between the application and FlashDreams are: + +- `ApplicationArguments`: runner mode and unparsed application arguments. +- `AppConfig`: presentation metadata consumed by runner I/O modes. +- `InferenceInput`: global conditioning at session creation and optional + per-step input. +- `StepRequest`: shared-serving request for the next iteration. +- `StepResult`: generated video chunk, layout, metadata, and metrics. +- `OutputArtifact`: persistent output returned by an I/O handler. + +Pipeline presets, WAN recipe classes, autoregressive cache contents, browser +routes, and MP4 recording are implementation details rather than part of the +minimal application ABI. + +## Ownership boundary + +- `flashdreams-runner` owns application discovery, CLI modes, device/process + setup, lifecycle, iteration, and output presentation. +- `flashdreams.runtime` owns shared inference inputs, requests, results, + artifacts, and serving protocols. +- `flashdreams.infra` owns reusable pipeline, decoder, post-processing, and + configuration primitives. +- `t2v_app` owns T2V arguments, presets, pipeline setup, session state, + generation, and optional WebRTC behavior. diff --git a/apps/t2v_app/README.md b/apps/t2v_app/README.md new file mode 100644 index 000000000..df99daa0d --- /dev/null +++ b/apps/t2v_app/README.md @@ -0,0 +1,61 @@ +# T2V Example Application + +`t2v-app` is an example implementation of the `flashdreams-runner` +application ABI. Its public module exports only `create_runtime(arguments)`. + +The implementation has three layers: + +- `application.py` parses T2V arguments, resolves the YAML pipeline preset, and + returns an uninitialized `T2VRuntime`. +- `runtime.py` owns pipeline construction, model weights, and one-time device + initialization. +- `session.py` owns the prompt, autoregressive cache, step counter, and the + pipeline `generate`/`finalize` calls for each main-loop iteration. + +The runner owns mode selection, process setup, session lifecycle, iteration, +and presentation. + +```bash +uv run flashdreams-runner t2v-app mp4 \ + --preset-id causal-forcing-wan2.1-t2v-1.3b-chunkwise \ + --prompt "A waterfall at sunset" \ + --output outputs/waterfall.mp4 + +uv run flashdreams-runner t2v-app webrtc \ + --preset-id self-forcing-wan2.1-t2v-1.3b + +uv run flashdreams-runner t2v-app none \ + --steps 2 \ + --prompt "A waterfall" +``` + +## Pipeline presets + +The application loads a YAML preset catalog through +`flashdreams.core.pipeline_presets` and asks the selected pipeline provider to +construct a `StreamInferencePipelineConfig`. The packaged catalog is +[`t2v_app/pipeline_presets.yaml`](t2v_app/pipeline_presets.yaml); pass +`--preset-config` to use another catalog. + +Every preset specifies a pipeline provider, application defaults, and +provider-owned pipeline options. `total_blocks` is an optional default for +finite runner modes; `--steps` overrides it. The T2V WebRTC page lets each user +edit the prompt and video duration, keeps the connection open for subsequent +generations, plays the completed MP4, and downloads a ZIP containing the video +and prompt metadata. + +FlashDreams' `ObjectGraphPipelineProvider` supports these trusted declarative +nodes: + +- `_target: module:attribute` imports and calls a config class with the other + mapping entries as keyword arguments. +- `_ref: module:attribute` imports a value such as a checkpoint transform + without calling it. +- `_tuple: [...]` preserves tuple-valued config fields. + +A custom package can expose a zero-argument pipeline provider class or instance +implementing `flashdreams.core.pipeline_presets.PipelineProvider` and reference +it from the YAML `provider` field. + +See the [`flashdreams-runner` application ABI](../../flashdreams_runner/README.md#application-abi) +for the runtime, session, and mode lifecycle. diff --git a/apps/t2v_app/pyproject.toml b/apps/t2v_app/pyproject.toml new file mode 100644 index 000000000..531e67e0a --- /dev/null +++ b/apps/t2v_app/pyproject.toml @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "t2v-app" +version = "0.1.0" +description = "Text-to-video example application for flashdreams-runner" +readme = "README.md" +requires-python = ">=3.10" +dependencies = ["flashdreams", "flashdreams-runner"] + +[tool.uv.sources] +flashdreams = { workspace = true } +flashdreams-runner = { workspace = true } + +[tool.pyright] +extraPaths = ["../..", "../../flashdreams"] +venvPath = "../.." +venv = ".venv" + +[tool.setuptools.packages.find] +where = ["."] + +[tool.setuptools.package-data] +t2v_app = ["pipeline_presets.yaml", "web/*.css", "web/*.js"] diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py new file mode 100644 index 000000000..90a014158 --- /dev/null +++ b/apps/t2v_app/t2v_app/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text-to-video example application for ``flashdreams-runner``.""" + +from .application import create_runtime + +__all__ = ["create_runtime"] diff --git a/apps/t2v_app/t2v_app/application.py b/apps/t2v_app/t2v_app/application.py new file mode 100644 index 000000000..9923e8c31 --- /dev/null +++ b/apps/t2v_app/t2v_app/application.py @@ -0,0 +1,201 @@ +# 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. + +"""Text-to-video application runtime factory.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path + +from flashdreams.core.pipeline_presets import load_pipeline_provider +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams_runner import AppConfig, ApplicationArguments, Runtime + +from .presets import PipelinePreset, RuntimePresetOptions, load_preset_catalog +from .runtime import T2VRuntime +from .session import T2VSessionDefaults + +FIELD_PROMPT = "prompt" +FIELD_TOTAL_BLOCKS = "total_blocks" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" +FIELD_FPS = "fps" + + +def create_runtime(arguments: ApplicationArguments) -> Runtime: + """Parse T2V options and create an uninitialized application runtime. + + Args: + arguments: Runner request containing the selected mode and parser. + + Returns: + Runtime containing resolved pipeline and session configuration. + """ + parser = arguments.parser + parser.add_argument( + "--preset-config", + type=Path, + help="Pipeline preset YAML (defaults to t2v_app's packaged catalog)", + ) + parser.add_argument( + "--preset-id", + help="Preset key (defaults to default_preset_id from the YAML)", + ) + parser.add_argument("--prompt") + parser.add_argument("--total-blocks", type=int) + parser.add_argument("--height", type=int, dest=FIELD_PIXEL_HEIGHT) + parser.add_argument("--width", type=int, dest=FIELD_PIXEL_WIDTH) + parser.add_argument("--fps", type=int) + options = vars(arguments.parse_args()) + + preset_id, preset = _resolve_preset(options) + scenario = _scenario(options, preset.runtime) + total_steps = _total_steps(options, preset.runtime) + return T2VRuntime( + pipeline_config=_create_pipeline_config(preset_id, preset), + session_defaults=T2VSessionDefaults( + prompt=str(scenario[FIELD_PROMPT]), + total_blocks=total_steps, + pixel_height=_required_int( + scenario[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT + ), + pixel_width=_required_int( + scenario[FIELD_PIXEL_WIDTH], name=FIELD_PIXEL_WIDTH + ), + fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), + ), + config=AppConfig( + model_id="t2v-app", + fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), + output_layout=preset.runtime.output_layout, + video_width=_required_int( + scenario[FIELD_PIXEL_WIDTH], name=FIELD_PIXEL_WIDTH + ), + video_height=_required_int( + scenario[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT + ), + default_steps=total_steps, + ), + ) + + +def _resolve_preset( + options: Mapping[str, object], +) -> tuple[str, PipelinePreset[RuntimePresetOptions]]: + """Resolve the configured pipeline preset.""" + catalog = load_preset_catalog(_optional_path(options.get("preset_config"))) + return catalog.resolve(_optional_string(options.get("preset_id"))) + + +def _create_pipeline_config( + preset_id: str, + preset: PipelinePreset[RuntimePresetOptions], +) -> StreamInferencePipelineConfig: + """Create the pipeline config selected by a resolved preset.""" + pipeline_provider = load_pipeline_provider(preset.provider) + pipeline_config = pipeline_provider.create_pipeline_config( + preset_id=preset_id, + options=preset.pipeline, + ) + if not isinstance(pipeline_config, StreamInferencePipelineConfig): + raise TypeError( + f"Pipeline provider returned {type(pipeline_config).__name__}, " + "expected StreamInferencePipelineConfig." + ) + if pipeline_config.name != preset_id: + raise ValueError( + f"Preset {preset_id!r} constructed pipeline " + f"{pipeline_config.name!r}; the preset key and pipeline name must match." + ) + return pipeline_config + + +def _scenario( + options: Mapping[str, object], + defaults: RuntimePresetOptions, +) -> dict[str, object]: + prompt_value = options.get(FIELD_PROMPT) + prompt = _resolve_prompt(defaults.prompt if prompt_value is None else prompt_value) + scenario = { + FIELD_PROMPT: prompt, + FIELD_PIXEL_HEIGHT: _option_or_default( + options, FIELD_PIXEL_HEIGHT, defaults.pixel_height + ), + FIELD_PIXEL_WIDTH: _option_or_default( + options, FIELD_PIXEL_WIDTH, defaults.pixel_width + ), + FIELD_FPS: _option_or_default(options, FIELD_FPS, defaults.fps), + } + for name in (FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH, FIELD_FPS): + if _required_int(scenario[name], name=name) <= 0: + raise ValueError(f"{name} must be > 0.") + return scenario + + +def _total_steps( + options: Mapping[str, object], + defaults: RuntimePresetOptions, +) -> int | None: + value = options.get(FIELD_TOTAL_BLOCKS) + if value is None: + value = defaults.total_blocks + if value is None: + return None + total_steps = _required_int(value, name=FIELD_TOTAL_BLOCKS) + if total_steps <= 0: + raise ValueError(f"{FIELD_TOTAL_BLOCKS} must be > 0.") + return total_steps + + +def _resolve_prompt(value: object) -> str: + if isinstance(value, Path): + lines = (line.strip() for line in value.read_text().splitlines()) + prompt = next((line for line in lines if line), "") + else: + prompt = str(value).strip() + if not prompt: + raise ValueError("A non-empty text-to-video prompt is required.") + return prompt + + +def _option_or_default( + options: Mapping[str, object], name: str, default: object +) -> object: + value = options.get(name) + return default if value is None else value + + +def _required_int(value: object, *, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + raise TypeError(f"{name} must be integer-compatible, got {value!r}.") + return int(value) + + +def _optional_path(value: object) -> str | Path | None: + if value is None or isinstance(value, (str, Path)): + return value + raise TypeError(f"Expected path or None, got {type(value).__name__}.") + + +def _optional_string(value: object) -> str | None: + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + raise TypeError("preset_id must be a non-empty string or None.") + return value.strip() + + +__all__ = ["create_runtime"] diff --git a/apps/t2v_app/t2v_app/pipeline_presets.yaml b/apps/t2v_app/t2v_app/pipeline_presets.yaml new file mode 100644 index 000000000..483c4620b --- /dev/null +++ b/apps/t2v_app/t2v_app/pipeline_presets.yaml @@ -0,0 +1,459 @@ +schema_version: 1 +default_preset_id: causal-forcing-wan2.1-t2v-1.3b-chunkwise + +presets: + causal-forcing-wan2.1-t2v-1.3b-chunkwise: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A cinematic waterfall cascading through a lush forest at sunset. + total_blocks: 60 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: causal-forcing-wan2.1-t2v-1.3b-chunkwise + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: true + cp_method: ring + checkpoint_path: https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/chunkwise/causal_forcing.pt + checkpoint_min_free_gb: null + state_dict_transform: + _ref: flashdreams.core.checkpoint.remap:unwrap_generator_state_dict + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 3 + guidance_scale: 1.0 + window_size_t: 21 + sink_size_t: 0 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: true + use_cuda_graph: true + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler.fm:FlowMatchSchedulerConfig + num_inference_steps: 4 + denoising_timesteps: [1000, 750, 500, 250] + warp_denoising_step: true + shift: 5.0 + sigma_max: 1.0 + sigma_min: 0.0 + extra_one_step: true + num_train_timesteps: 1000 + enable_tqdm: false + + causal-forcing-wan2.1-t2v-1.3b-framewise: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A cinematic waterfall cascading through a lush forest at sunset. + total_blocks: 60 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: causal-forcing-wan2.1-t2v-1.3b-framewise + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: true + cp_method: ring + checkpoint_path: https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/framewise/causal_forcing.pt + checkpoint_min_free_gb: null + state_dict_transform: + _ref: flashdreams.core.checkpoint.remap:unwrap_generator_state_dict + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 1 + guidance_scale: 1.0 + window_size_t: 21 + sink_size_t: 0 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: true + use_cuda_graph: true + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler.fm:FlowMatchSchedulerConfig + num_inference_steps: 4 + denoising_timesteps: [1000, 750, 500, 250] + warp_denoising_step: true + shift: 5.0 + sigma_max: 1.0 + sigma_min: 0.0 + extra_one_step: true + num_train_timesteps: 1000 + enable_tqdm: false + + self-forcing-wan2.1-t2v-1.3b: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A woman walks through a neon-lit city street at night. + total_blocks: 60 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: self-forcing-wan2.1-t2v-1.3b + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: true + cp_method: ring + checkpoint_path: https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt + checkpoint_min_free_gb: null + state_dict_transform: + _ref: flashdreams.core.checkpoint.remap:unwrap_generator_state_dict + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 3 + guidance_scale: 1.0 + window_size_t: 21 + sink_size_t: 0 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: true + use_cuda_graph: true + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler.fm:FlowMatchSchedulerConfig + num_inference_steps: 4 + denoising_timesteps: [1000, 750, 500, 250] + warp_denoising_step: true + shift: 8.0 + sigma_max: 1.0 + sigma_min: 0.0 + extra_one_step: true + num_train_timesteps: 1000 + enable_tqdm: false + + self-forcing-wan2.1-t2v-1.3b-taehv: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A woman walks through a neon-lit city street at night. + total_blocks: 60 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: self-forcing-wan2.1-t2v-1.3b-taehv + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.taehv:TeahvVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/lighttaew2_1.pth + state_dict_transform: + _ref: flashdreams.recipes.taehv:lighttae_state_dict_transform + use_cuda_graph: true + use_compile: true + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: true + cp_method: ring + checkpoint_path: https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt + checkpoint_min_free_gb: null + state_dict_transform: + _ref: flashdreams.core.checkpoint.remap:unwrap_generator_state_dict + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 3 + guidance_scale: 1.0 + window_size_t: 21 + sink_size_t: 0 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: true + use_cuda_graph: true + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler.fm:FlowMatchSchedulerConfig + num_inference_steps: 4 + denoising_timesteps: [1000, 750, 500, 250] + warp_denoising_step: true + shift: 8.0 + sigma_max: 1.0 + sigma_min: 0.0 + extra_one_step: true + num_train_timesteps: 1000 + enable_tqdm: false + + self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A woman walks through a neon-lit city street at night. + total_blocks: 80 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: self-forcing-wan2.1-t2v-1.3b-sink5-window7-rerope + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 0 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: false + cp_method: ring + checkpoint_path: https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt + checkpoint_min_free_gb: null + state_dict_transform: + _ref: flashdreams.core.checkpoint.remap:unwrap_generator_state_dict + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 3 + guidance_scale: 1.0 + window_size_t: 7 + sink_size_t: 5 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: false + use_cuda_graph: false + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler.fm:FlowMatchSchedulerConfig + num_inference_steps: 4 + denoising_timesteps: [1000, 750, 500, 250] + warp_denoising_step: true + shift: 8.0 + sigma_max: 1.0 + sigma_min: 0.0 + extra_one_step: true + num_train_timesteps: 1000 + enable_tqdm: false + + wan21-t2v-1.3b-480p: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A white cat wearing sunglasses sits on a surfboard at the beach. + total_blocks: 1 + pixel_height: 480 + pixel_width: 832 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.wan:WanInferencePipelineConfig + name: wan21-t2v-1.3b-480p + enable_sync_and_profile: true + encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.umt5:UMT5TextEncoderConfig + model_id_or_local_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + image_encoder: null + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.wan:Wan21TransformerConfig + network: + _target: flashdreams.recipes.wan:WanDiTNetwork1pt3BConfig + patch_embedding_type: conv3d + apply_rope_before_kvcache: true + cp_method: ring + checkpoint_path: https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B/blob/main/diffusion_pytorch_model.safetensors + checkpoint_min_free_gb: null + state_dict_transform: null + stream_checkpoint: false + init_device: null + batch_shape: + _tuple: [] + len_t: 21 + guidance_scale: 6.0 + window_size_t: 21 + sink_size_t: 0 + h_extrapolation_ratio: 1.0 + w_extrapolation_ratio: 1.0 + compile_network: true + use_cuda_graph: true + cuda_graph_warmup_iters: 2 + stamp_image_latent: false + concat_image_mask_to_latent: false + scheduler: + _target: flashdreams.infra.diffusion.scheduler:FlowMatchUniPCSchedulerConfig + num_inference_steps: 50 + shift: 8.0 + num_train_timesteps: 1000 + solver_order: 2 + use_kerras_sigma: false + enable_tqdm: true + + cosmos2-t2v-2b-720p: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: A robotic arm performs precision welding in an industrial workshop. + total_blocks: 1 + pixel_height: 720 + pixel_width: 1280 + fps: 16 + output_layout: tchw + pipeline: + _target: flashdreams.recipes.cosmos.pipeline:CosmosInferencePipelineConfig + name: cosmos2-t2v-2b-720p + enable_sync_and_profile: true + encoder: null + image_encoder: null + text_encoder: + _target: flashdreams.infra.encoder.text.cosmos_reason1:CosmosReason1TextEncoderConfig + model_name: nvidia/Cosmos-Reason1-7B + revision: 3210bec0495fdc7a8d3dbb8d58da5711eab4b423 + max_length: 512 + embedding_concat_strategy: full_concat + n_layers_per_group: 5 + decoder: + _target: flashdreams.recipes.wan:WanVAEDecoderConfig + checkpoint_path: https://huggingface.co/lightx2v/Autoencoders/resolve/main/Wan2.1_VAE.pth + use_cuda_graph: true + use_compile: false + diffusion_model: + _target: flashdreams.infra.diffusion.model:DiffusionModelConfig + seed: 42 + context_noise: 0 + noise_in_unpatchified_shape: false + transformer: + _target: flashdreams.recipes.cosmos.transformer:CosmosTransformerConfig + network: + _target: flashdreams.recipes.cosmos.transformer.impl.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: + _ref: flashdreams.recipes.cosmos.transformer.impl.network:state_dict_transform + batch_shape: + _tuple: [] + len_t: 24 + h_extrapolation_ratio: 3.0 + w_extrapolation_ratio: 3.0 + window_size_t: 24 + sink_size_t: 0 + compile_network: true + use_cuda_graph: false + cuda_graph_warmup_iters: 2 + skip_finalize_kv_cache: false + guidance_scale: 8.0 + conditional_frame_timestep: null + scheduler: + _target: flashdreams.infra.diffusion.scheduler:FlowMatchUniPCSchedulerConfig + num_inference_steps: 35 + shift: 5.0 + num_train_timesteps: 1000 + solver_order: 2 + use_kerras_sigma: true + enable_tqdm: true diff --git a/apps/t2v_app/t2v_app/presets.py b/apps/t2v_app/t2v_app/presets.py new file mode 100644 index 000000000..2c9e8ea1e --- /dev/null +++ b/apps/t2v_app/t2v_app/presets.py @@ -0,0 +1,170 @@ +# 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. + +"""Text-to-video runtime options for shared pipeline-preset catalogs.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from importlib.resources import files +from pathlib import Path +from typing import Any, cast + +from flashdreams.core.pipeline_presets import ( + PipelinePreset, + PresetCatalog, + load_pipeline_preset_catalog, + parse_pipeline_preset_catalog, +) +from flashdreams.infra.postprocess import VideoTensorLayout + +_REQUIRED_RUNTIME_FIELDS = { + "prompt", + "pixel_height", + "pixel_width", + "fps", + "output_layout", +} +_OPTIONAL_RUNTIME_FIELDS = {"total_blocks"} +_VIDEO_LAYOUTS = {"tchw", "btchw", "bcthw", "bvtchw"} + + +@dataclass(frozen=True, slots=True) +class RuntimePresetOptions: + """Host-facing rollout and presentation options for one T2V preset.""" + + prompt: str + """Default text prompt.""" + + pixel_height: int + """Output video height in pixels.""" + + pixel_width: int + """Output video width in pixels.""" + + fps: int + """Presentation frame rate.""" + + output_layout: VideoTensorLayout + """Decoded tensor layout exposed to the host.""" + + total_blocks: int | None = None + """Default finite-session length; ``None`` requires an MP4 CLI override.""" + + +def load_preset_catalog( + path: str | Path | None = None, +) -> PresetCatalog[RuntimePresetOptions]: + """Load the packaged or caller-supplied T2V pipeline-preset catalog. + + Args: + path: YAML path; ``None`` loads the catalog packaged with ``t2v_app``. + + Returns: + Validated preset catalog. + """ + if path is not None: + return load_pipeline_preset_catalog( + path, + runtime_options_parser=_load_runtime_options, + ) + + source = files("t2v_app").joinpath("pipeline_presets.yaml") + return parse_pipeline_preset_catalog( + source.read_text(encoding="utf-8"), + source_name=str(source), + runtime_options_parser=_load_runtime_options, + ) + + +def _load_runtime_options(value: object, *, path: str) -> RuntimePresetOptions: + runtime = _mapping(value, path=path) + _require_fields( + runtime, + required=_REQUIRED_RUNTIME_FIELDS, + optional=_OPTIONAL_RUNTIME_FIELDS, + path=path, + ) + layout = _nonempty_string(runtime["output_layout"], path=f"{path}.output_layout") + if layout not in _VIDEO_LAYOUTS: + allowed = ", ".join(sorted(_VIDEO_LAYOUTS)) + raise ValueError( + f"{path}.output_layout must be one of {allowed}, got {layout!r}." + ) + return RuntimePresetOptions( + prompt=_nonempty_string(runtime["prompt"], path=f"{path}.prompt"), + pixel_height=_positive_int( + runtime["pixel_height"], path=f"{path}.pixel_height" + ), + pixel_width=_positive_int(runtime["pixel_width"], path=f"{path}.pixel_width"), + fps=_positive_int(runtime["fps"], path=f"{path}.fps"), + output_layout=cast(VideoTensorLayout, layout), + total_blocks=( + None + if runtime.get("total_blocks") is None + else _positive_int(runtime["total_blocks"], path=f"{path}.total_blocks") + ), + ) + + +def _mapping(value: object, *, path: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise TypeError(f"{path} must be a mapping, got {type(value).__name__}.") + mapping = cast(dict[object, object], dict(cast(Any, value))) + if any(not isinstance(key, str) for key in mapping): + raise TypeError(f"{path} keys must be strings.") + return cast(dict[str, object], mapping) + + +def _require_fields( + value: Mapping[str, object], + *, + required: set[str], + optional: set[str], + path: str, +) -> None: + fields = {str(key) for key in value} + missing = required - fields + unknown = fields - required - optional + if missing or unknown: + details: list[str] = [] + if missing: + details.append(f"missing {sorted(missing)}") + if unknown: + details.append(f"unknown {sorted(unknown)}") + raise ValueError(f"Invalid fields at {path}: {'; '.join(details)}.") + + +def _nonempty_string(value: object, *, path: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise TypeError(f"{path} must be a non-empty string.") + return value.strip() + + +def _positive_int(value: object, *, path: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{path} must be an integer.") + if value <= 0: + raise ValueError(f"{path} must be > 0.") + return value + + +__all__ = [ + "PipelinePreset", + "PresetCatalog", + "RuntimePresetOptions", + "load_preset_catalog", +] diff --git a/apps/t2v_app/t2v_app/runtime.py b/apps/t2v_app/t2v_app/runtime.py new file mode 100644 index 000000000..53edaa73d --- /dev/null +++ b/apps/t2v_app/t2v_app/runtime.py @@ -0,0 +1,172 @@ +# 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. + +"""Text-to-video model runtime and one-time pipeline state.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast + +import torch + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineConfig, +) +from flashdreams.runtime import InferenceInput +from flashdreams_runner import AppConfig, IOHandler, Runtime, Session +from flashdreams_runner.webrtc import WebRTCMode + +from .session import T2VScenario, T2VSession, T2VSessionDefaults + + +@dataclass(frozen=True, slots=True) +class T2VArtifact: + """Completed WebRTC recording and the scenario that produced it.""" + + path: Path + """Path to the generated MP4 file.""" + + scenario: T2VScenario + """Prompt, duration, and video geometry stored with the recording.""" + + +class T2VRuntime(Runtime): + """Own T2V model weights and create isolated generation sessions.""" + + def __init__( + self, + *, + pipeline_config: StreamInferencePipelineConfig, + session_defaults: T2VSessionDefaults, + config: AppConfig, + ) -> None: + self._pipeline_config = pipeline_config + self._session_defaults = session_defaults + self._config = config + self._pipeline: StreamInferencePipeline[Any, Any, Any] | None = None + self._io_handler: IOHandler | None = None + self._record_sessions = False + self._latest_artifact: T2VArtifact | None = None + + @property + def config(self) -> AppConfig: + """Return T2V configuration for runner-owned presentation.""" + return self._config + + def initialize(self, *, device: str, io_handler: IOHandler) -> None: + """Construct model weights once for the selected device and I/O mode.""" + if self._pipeline is not None: + raise RuntimeError("T2VRuntime is already initialized.") + pipeline = self._pipeline_config.setup() + if not isinstance(pipeline, StreamInferencePipeline): + raise TypeError( + "T2V pipeline config must construct StreamInferencePipeline, got " + f"{type(pipeline).__name__}." + ) + self._pipeline = pipeline.to(device).eval() + self._io_handler = io_handler + if isinstance(io_handler, WebRTCMode): + from .webrtc import T2VWebRTCCustomization + + self._record_sessions = True + io_handler.customize(T2VWebRTCCustomization(runtime=self)) + + def create_session(self, initial_input: InferenceInput | None = None) -> Session: + """Create a T2V session with its own prompt and autoregressive cache.""" + if self._pipeline is None: + raise RuntimeError("T2VRuntime must be initialized before use.") + return T2VSession( + pipeline=self._pipeline, + defaults=self._session_defaults, + initial_input=initial_input or InferenceInput(), + output_layout=self._config.output_layout, + record_artifact=self._record_artifact if self._record_sessions else None, + ) + + def prepare_session_input( + self, + *, + prompt: str | None = None, + total_blocks: int | None = None, + ) -> InferenceInput: + """Build complete initial input for a browser-created T2V session.""" + return InferenceInput( + global_conditioning={ + "prompt": self._session_defaults.prompt if prompt is None else prompt, + "total_blocks": ( + self._session_defaults.total_blocks + if total_blocks is None + else total_blocks + ), + "pixel_height": self._session_defaults.pixel_height, + "pixel_width": self._session_defaults.pixel_width, + "fps": self._session_defaults.fps, + } + ) + + def blocks_for_duration(self, duration_s: float) -> int: + """Return enough autoregressive blocks for a requested duration.""" + if not math.isfinite(duration_s) or duration_s <= 0: + raise ValueError("duration_s must be finite and > 0.") + pipeline = self._pipeline + if pipeline is None: + raise RuntimeError("T2VRuntime must be initialized before use.") + target_frames = math.ceil(duration_s * self._session_defaults.fps) + generated_frames = 0 + block_index = 0 + pipeline_api = cast(Any, pipeline) + while generated_frames < target_frames: + block_frames = int(pipeline_api.get_num_output_frames(block_index)) + if block_frames <= 0: + raise ValueError("T2V pipeline output frame counts must be > 0.") + generated_frames += block_frames + block_index += 1 + return block_index + + def peek_steady_output_num_frames(self) -> int: + """Return the steady chunk size used to bound WebRTC delivery queues.""" + pipeline = self._pipeline + if pipeline is None: + raise RuntimeError("T2VRuntime must be initialized before use.") + return int(cast(Any, pipeline).get_num_output_frames(1)) + + @property + def latest_artifact(self) -> T2VArtifact | None: + """Return the most recently completed WebRTC recording.""" + return self._latest_artifact + + def _record_artifact(self, path: Path, scenario: T2VScenario) -> None: + self._latest_artifact = T2VArtifact(path=path, scenario=scenario) + + def destroy(self) -> None: + """Release pipeline weights and accelerator allocator state.""" + pipeline = self._pipeline + self._pipeline = None + self._io_handler = None + self._record_sessions = False + if pipeline is None: + return + close = getattr(pipeline, "close", None) + if callable(close): + close() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +__all__ = ["T2VArtifact", "T2VRuntime"] diff --git a/apps/t2v_app/t2v_app/session.py b/apps/t2v_app/t2v_app/session.py new file mode 100644 index 000000000..901856b09 --- /dev/null +++ b/apps/t2v_app/t2v_app/session.py @@ -0,0 +1,287 @@ +# 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. + +"""Text-to-video session state and generation loop iteration.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast +from uuid import uuid4 + +import torch + +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.pipeline import StreamInferencePipeline +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.runtime import InferenceInput, StepRequest, StepResult +from flashdreams.runtime.video_output import Mp4VideoOutputTarget +from flashdreams_runner import Session + +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: + """Validated state used by one text-to-video session.""" + + prompt: str + """Text prompt used to initialize the model cache.""" + + total_blocks: int | None + """Generation limit, or ``None`` for an externally driven session.""" + + pixel_height: int + """Output height in pixels.""" + + pixel_width: int + """Output width in pixels.""" + + fps: int + """Output video frame rate.""" + + +@dataclass(frozen=True, kw_only=True, slots=True) +class T2VSessionDefaults: + """Default state copied into each new T2V session.""" + + prompt: str + """Text prompt used when the session input does not override it.""" + + total_blocks: int | None + """Optional session-owned generation limit.""" + + pixel_height: int + """Output height used when the session input does not override it.""" + + pixel_width: int + """Output width used when the session input does not override it.""" + + fps: int + """Output frame rate used for WebRTC recordings.""" + + +class T2VSession(Session): + """Own one prompt, autoregressive cache, and T2V generation loop.""" + + def __init__( + self, + *, + pipeline: StreamInferencePipeline[Any, Any, Any], + defaults: T2VSessionDefaults, + initial_input: InferenceInput, + output_layout: VideoTensorLayout, + record_artifact: Callable[[Path, T2VScenario], None] | None = None, + recording_directory: Path | None = None, + ) -> None: + scenario = _session_scenario(defaults, initial_input) + decoder = pipeline.decoder + if not isinstance(decoder, StreamingVideoDecoder): + raise TypeError("T2V pipelines require a StreamingVideoDecoder.") + ratio = decoder.spatial_compression_ratio + pixel_height = scenario.pixel_height + pixel_width = scenario.pixel_width + if pixel_height % ratio or pixel_width % ratio: + raise ValueError( + "T2V dimensions must be divisible by the decoder spatial " + f"compression ratio ({ratio})." + ) + + self._pipeline = pipeline + self._scenario = scenario + self._prompt = scenario.prompt + self._pixel_height = pixel_height + self._pixel_width = pixel_width + self._output_layout: VideoTensorLayout = output_layout + pipeline_api = cast(Any, pipeline) + self._cache: object | None = pipeline_api.initialize_cache( + text=[self._prompt], + image=None, + height=pixel_height // ratio, + width=pixel_width // ratio, + ) + self._step_index = 0 + self._steady_output_frame_count = int(pipeline_api.get_num_output_frames(1)) + self._destroyed = False + self._record_artifact = record_artifact + self._artifact_path: Path | None = None + self._artifact_output: Mp4VideoOutputTarget | None = None + if record_artifact is not None: + recording_directory = recording_directory or Path("outputs/t2v-webrtc") + recording_directory.mkdir(parents=True, exist_ok=True) + self._artifact_path = recording_directory / f"{uuid4()}.mp4" + self._artifact_output = Mp4VideoOutputTarget( + output_path=self._artifact_path, + fps=scenario.fps, + output_layout=output_layout, + ) + self._artifact_output.open() + + @property + def step_index(self) -> int: + """Return the index of the next autoregressive block.""" + return self._step_index + + @property + def steady_output_frame_count(self) -> int: + """Return the steady decoded frames produced by one iteration.""" + return self._steady_output_frame_count + + def next_step_request(self) -> StepRequest | None: + """Stop shared serving after this session's requested video duration.""" + total_blocks = self._scenario.total_blocks + if self._destroyed or ( + total_blocks is not None and self._step_index >= total_blocks + ): + return None + return StepRequest( + step_index=self._step_index, + metadata={ + "steady_output_frame_count": self._steady_output_frame_count, + }, + ) + + def step(self, inputs: InferenceInput) -> StepResult: + """Generate and finalize one autoregressive video block.""" + if self._destroyed or self._cache is None: + raise RuntimeError("Cannot generate from a destroyed T2VSession.") + if inputs.global_conditioning: + raise ValueError( + "T2V global conditioning is fixed when the session is created." + ) + if inputs.step: + raise ValueError("This T2V application does not accept per-step input.") + + index = self._step_index + video = self._pipeline.generate( + autoregressive_index=index, + cache=cast(Any, self._cache), + ) + metrics = _metrics( + self._pipeline.finalize( + autoregressive_index=index, + cache=cast(Any, self._cache), + ) + ) + self._step_index += 1 + if not isinstance(video, torch.Tensor): + raise TypeError( + "T2V pipeline generate() must return torch.Tensor, got " + f"{type(video).__name__}." + ) + result = StepResult.from_video_chunk( + step_index=index, + video_chunk=video.detach(), + layout=self._output_layout, + metadata={FIELD_PROMPT: self._prompt}, + metrics=metrics, + ) + if self._artifact_output is not None: + self._artifact_output.write(result) + return result + + def destroy(self) -> None: + """Release this session's autoregressive cache.""" + if self._destroyed: + return + self._destroyed = True + cache = self._cache + self._cache = None + artifact_output = self._artifact_output + artifact_path = self._artifact_path + self._artifact_output = None + self._artifact_path = None + try: + if artifact_output is not None: + artifacts = artifact_output.close() + if ( + artifacts + and artifact_path is not None + and self._record_artifact is not None + ): + self._record_artifact(artifact_path, self._scenario) + finally: + close = getattr(cache, "close", None) + if callable(close): + close() + + +def _session_scenario( + defaults: T2VSessionDefaults, + initial_input: InferenceInput, +) -> T2VScenario: + values = { + FIELD_PROMPT: defaults.prompt, + FIELD_TOTAL_BLOCKS: defaults.total_blocks, + FIELD_PIXEL_HEIGHT: defaults.pixel_height, + FIELD_PIXEL_WIDTH: defaults.pixel_width, + FIELD_FPS: defaults.fps, + } + values.update(initial_input.global_conditioning) + prompt = str(values[FIELD_PROMPT]).strip() + if not prompt: + raise ValueError("A non-empty text-to-video prompt is required.") + values[FIELD_PROMPT] = prompt + pixel_height = _positive_int(values[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT) + pixel_width = _positive_int(values[FIELD_PIXEL_WIDTH], name=FIELD_PIXEL_WIDTH) + fps = _positive_int(values[FIELD_FPS], name=FIELD_FPS) + total_blocks = values[FIELD_TOTAL_BLOCKS] + if total_blocks is not None: + if isinstance(total_blocks, bool) or not isinstance(total_blocks, int): + raise TypeError(f"{FIELD_TOTAL_BLOCKS} must be an integer or None.") + if total_blocks <= 0: + raise ValueError(f"{FIELD_TOTAL_BLOCKS} must be > 0.") + return T2VScenario( + prompt=prompt, + total_blocks=total_blocks, + pixel_height=pixel_height, + pixel_width=pixel_width, + fps=fps, + ) + + +def _positive_int(value: object, *, name: str) -> int: + 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.") + return value + + +def _metrics(value: object) -> Mapping[str, float | int]: + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError( + "T2V pipeline finalize() must return a metrics mapping or None, got " + f"{type(value).__name__}." + ) + metrics: dict[str, float | int] = {} + for key, metric in value.items(): + if not isinstance(key, str): + raise TypeError("T2V pipeline metric keys must be strings.") + if not isinstance(metric, (int, float)): + raise TypeError(f"T2V pipeline metric {key!r} must be numeric.") + metrics[key] = metric + return metrics + + +__all__ = ["T2VScenario", "T2VSession", "T2VSessionDefaults"] diff --git a/apps/t2v_app/t2v_app/web/adapter.css b/apps/t2v_app/t2v_app/web/adapter.css new file mode 100644 index 000000000..7956dc8d5 --- /dev/null +++ b/apps/t2v_app/t2v_app/web/adapter.css @@ -0,0 +1,10 @@ +/* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. */ +/* SPDX-License-Identifier: Apache-2.0 */ + +.promptGenerationPanel textarea { + resize: vertical; +} + +.promptDurationInput { + width: 7rem; +} diff --git a/apps/t2v_app/t2v_app/web/adapter.js b/apps/t2v_app/t2v_app/web/adapter.js new file mode 100644 index 000000000..f93392187 --- /dev/null +++ b/apps/t2v_app/t2v_app/web/adapter.js @@ -0,0 +1,31 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** T2V metadata and finite-generation controls for the shared WebRTC UI. */ +export default { + modelName: "Text-to-Video", + async mount(context) { + const response = await fetch("/api/t2v/config") + if (!response.ok) return + const config = await response.json() + context.setModelName(config.model_id || "Text-to-Video") + const panel = document.querySelector(".promptGenerationPanel") + const prompt = panel?.querySelector("textarea") + const duration = panel?.querySelector(".promptDurationInput") + if (prompt && typeof config.default_prompt === "string") { + prompt.value = config.default_prompt + } + if (duration && Number.isFinite(Number(config.default_duration_s))) { + duration.value = String(config.default_duration_s) + } + }, + promptGeneration: { + endpoint: "/api/t2v/prompt", + label: "Describe the video", + placeholder: "A cinematic drone shot over snowy mountains at sunrise", + generateLabel: "Generate video", + downloadEndpoint: "/api/t2v/download", + playbackEndpoint: "/api/t2v/playback", + hideControls: true, + }, +} diff --git a/apps/t2v_app/t2v_app/webrtc.py b/apps/t2v_app/t2v_app/webrtc.py new file mode 100644 index 000000000..088f0f646 --- /dev/null +++ b/apps/t2v_app/t2v_app/webrtc.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""T2V-specific WebRTC controls, recording, playback, and downloads.""" + +from __future__ import annotations + +import io +import json +import zipfile +from importlib.resources import files +from pathlib import Path +from typing import Any, Protocol, cast + +from aiohttp import web + +from flashdreams.runtime import InferenceInput, InferenceRuntime +from flashdreams.runtime.demo import ( + DemoSpec, + PreparedScenario, + RuntimeHost, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig +from flashdreams_runner import AppConfig, Runtime +from flashdreams_runner.webrtc import ModelInputProviderFactory + +DEFAULT_DURATION_S = 5.0 +"""Initial browser duration when a preset has no finite block count.""" + +MAX_DURATION_S = 60.0 +"""Prototype UI limit that bounds one browser generation request.""" + + +class _Scenario(Protocol): + @property + def prompt(self) -> str: ... + + @property + def total_blocks(self) -> int | None: ... + + @property + def pixel_height(self) -> int: ... + + @property + def pixel_width(self) -> int: ... + + @property + def fps(self) -> int: ... + + +class _Artifact(Protocol): + @property + def path(self) -> Path: ... + + @property + def scenario(self) -> _Scenario: ... + + +class _T2VRuntime(Protocol): + @property + def config(self) -> AppConfig: ... + + @property + def latest_artifact(self) -> _Artifact | None: ... + + def prepare_session_input( + self, + *, + prompt: str | None = None, + total_blocks: int | None = None, + ) -> InferenceInput: ... + + def blocks_for_duration(self, duration_s: float) -> int: ... + + +class T2VWebRTCSessionManager( + BaseWebRTCSessionManager[_T2VRuntime, WebRTCRuntimeConfig] +): + """Keep one browser connection alive across finite T2V generations.""" + + def update_generation(self, *, prompt: str, duration_s: float) -> None: + """Prepare the prompt and duration used by the next generation.""" + prompt = prompt.strip() + if not prompt: + raise ValueError("Prompt must be non-empty.") + if not 0 < duration_s <= MAX_DURATION_S: + raise ValueError( + f"Duration must be greater than 0 and at most {MAX_DURATION_S:g} " + "seconds." + ) + self._shared_scenario = PreparedScenario( + initial_inputs=self.runtime.prepare_session_input( + prompt=prompt, + total_blocks=self.runtime.blocks_for_duration(duration_s), + ) + ) + + +class T2VWebRTCCustomization: + """Install T2V browser assets and HTTP routes into the runner mode.""" + + def __init__(self, *, runtime: _T2VRuntime) -> None: + self._runtime = runtime + + def prepare_initial_input(self) -> InferenceInput: + """Return a complete finite input for the first browser session.""" + total_blocks = self._runtime.config.default_steps + if total_blocks is None: + total_blocks = self._runtime.blocks_for_duration(DEFAULT_DURATION_S) + return self._runtime.prepare_session_input(total_blocks=total_blocks) + + def create_session_manager( + self, + *, + runtime: Runtime, + output: WebRTCOutputSpec, + spec: DemoSpec, + scenario: PreparedScenario, + input_provider_factory: ModelInputProviderFactory, + ) -> BaseWebRTCSessionManager[Any, Any]: + """Create a finite-generation manager over the initialized runtime.""" + if runtime is not cast(object, self._runtime): + raise ValueError("T2V WebRTC customization received a different runtime.") + inference_runtime = cast(InferenceRuntime, runtime) + return T2VWebRTCSessionManager( + runtime=self._runtime, + runtime_config=cast(WebRTCRuntimeConfig, cast(object, output)), + fps=int(self._runtime.config.fps), + identity=self._runtime.config.model_id, + supported_control_keys=frozenset(), + shared_host=RuntimeHost(inference_runtime), + shared_spec=spec, + shared_scenario=scenario, + shared_model_input_provider_factory=input_provider_factory, + client_liveness_timeout_s=output.client_liveness_timeout_s, + keep_connection_after_completed=True, + runtime_ready=True, + ) + + def create_app_resources( + self, + *, + session_manager: BaseWebRTCSessionManager[Any, Any], + ) -> WebRTCAppResources: + """Return the T2V adapter assets and application-owned routes.""" + if not isinstance(session_manager, T2VWebRTCSessionManager): + raise TypeError("T2V WebRTC requires T2VWebRTCSessionManager.") + return WebRTCAppResources( + model_web_resource=files("t2v_app").joinpath("web"), + configure_app=lambda app: _configure_app( + app, + manager=session_manager, + ), + preload_name="FlashDreams T2V", + ) + + +def _configure_app( + app: web.Application, + *, + manager: T2VWebRTCSessionManager, +) -> None: + """Register T2V metadata, prompt, playback, and download endpoints.""" + + async def app_config(_: web.Request) -> web.StreamResponse: + config = manager.runtime.config + initial_input = manager.runtime.prepare_session_input() + return web.json_response( + { + "model_id": config.model_id, + "default_prompt": initial_input.global_conditioning["prompt"], + "default_duration_s": DEFAULT_DURATION_S, + } + ) + + async def update_generation(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 isinstance(duration_s, bool) or not isinstance(duration_s, int | float): + raise web.HTTPBadRequest(reason="Expected numeric duration_s.") + try: + manager.update_generation( + 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.") + if not artifact.path.is_file(): + raise web.HTTPNotFound(reason="Generated MP4 is no longer available.") + scenario = artifact.scenario + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + archive.write(artifact.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.path.is_file(): + raise web.HTTPNotFound(reason="No completed MP4 is available yet.") + return web.FileResponse(artifact.path) + + app.router.add_get("/api/t2v/config", app_config) + app.router.add_post("/api/t2v/prompt", update_generation) + app.router.add_get("/api/t2v/download", download) + app.router.add_get("/api/t2v/playback", playback) + + +__all__ = ["T2VWebRTCCustomization", "T2VWebRTCSessionManager"] diff --git a/apps/t2v_app/tests/test_application.py b/apps/t2v_app/tests/test_application.py new file mode 100644 index 000000000..cfef9418f --- /dev/null +++ b/apps/t2v_app/tests/test_application.py @@ -0,0 +1,254 @@ +# 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. + +"""CPU tests for the T2V application runtime and session boundary.""" + +from __future__ import annotations + +import argparse +from typing import Any, cast + +import pytest +import torch +import torch.nn as nn +import t2v_app +from t2v_app import application +from t2v_app.presets import PipelinePreset, PresetCatalog, RuntimePresetOptions +from t2v_app.runtime import T2VRuntime +from t2v_app.session import T2VSession + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineConfig, +) +from flashdreams.runtime import InferenceInput +from flashdreams_runner import ( + Application, + ApplicationArguments, + DriveSession, + IOHandler, + Runtime, +) + +pytestmark = pytest.mark.ci_cpu + + +def _preset() -> tuple[ + str, + PipelinePreset[RuntimePresetOptions], + PresetCatalog[RuntimePresetOptions], +]: + preset_id = "test-t2v" + preset = PipelinePreset( + provider="tests:provider", + runtime=RuntimePresetOptions( + prompt="default prompt", + total_blocks=2, + pixel_height=64, + pixel_width=96, + fps=12, + output_layout="tchw", + ), + pipeline={}, + ) + return ( + preset_id, + preset, + PresetCatalog(default_preset_id=preset_id, presets={preset_id: preset}), + ) + + +def test_application_module_conforms_to_single_factory_abi() -> None: + assert isinstance(t2v_app, Application) + assert cast(Any, t2v_app).__all__ == ["create_runtime"] + + +def test_create_runtime_parses_options_without_constructing_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline_constructed = False + + class Pipeline(StreamInferencePipeline[Any, Any, Any]): + def __init__(self, config: object) -> None: + nn.Module.__init__(self) + del config + nonlocal pipeline_constructed + pipeline_constructed = True + + preset_id, _, catalog = _preset() + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, Pipeline), + name=preset_id, + diffusion_model=cast(Any, None), + ) + + class PipelineProvider: + def create_pipeline_config( + self, *, preset_id: str, options: object + ) -> StreamInferencePipelineConfig: + assert preset_id == "test-t2v" + assert options == {} + return pipeline_config + + monkeypatch.setattr(application, "load_preset_catalog", lambda _: catalog) + monkeypatch.setattr( + application, "load_pipeline_provider", lambda _: PipelineProvider() + ) + arguments = ApplicationArguments( + mode="webrtc", + parser=argparse.ArgumentParser(), + argv=("--prompt", "A waterfall"), + ) + + runtime = application.create_runtime(arguments) + + assert isinstance(runtime, T2VRuntime) + assert runtime.config.video_width == 96 + assert runtime.config.fps == 12 + assert runtime.config.default_steps == 2 + assert arguments.options.prompt == "A waterfall" + assert not pipeline_constructed + + +def test_runtime_owns_pipeline_and_session_owns_generation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + class Cache: + def close(self) -> None: + calls.append("cache.close") + + class Decoder: + spatial_compression_ratio = 8 + + monkeypatch.setattr("t2v_app.session.StreamingVideoDecoder", Decoder) + + class Pipeline(StreamInferencePipeline[Any, Any, Any]): + def __init__(self, config: object) -> None: + nn.Module.__init__(self) + del config + self.decoder = cast(Any, Decoder()) + calls.append("pipeline.init") + + def to(self, device: object) -> "Pipeline": + assert str(device) == "cpu" + calls.append("pipeline.to") + return self + + def eval(self) -> "Pipeline": + calls.append("pipeline.eval") + return self + + def initialize_cache(self, **kwargs: object) -> Cache: + assert kwargs == { + "text": ["A waterfall"], + "image": None, + "height": 8, + "width": 12, + } + calls.append("session.initialize_cache") + return Cache() + + def get_num_output_frames(self, autoregressive_index: int) -> int: + assert autoregressive_index == 1 + return 3 + + def generate(self, autoregressive_index: int, cache: object) -> torch.Tensor: + assert autoregressive_index == 0 + assert isinstance(cache, Cache) + calls.append("session.pipeline.generate") + return torch.zeros((3, 3, 2, 2)) + + def finalize( + self, autoregressive_index: int, cache: object + ) -> dict[str, float]: + assert autoregressive_index == 0 + assert isinstance(cache, Cache) + calls.append("session.pipeline.finalize") + return {"step_ms": 1.0} + + def close(self) -> None: + calls.append("pipeline.close") + + preset_id, _, catalog = _preset() + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, Pipeline), + name=preset_id, + diffusion_model=cast(Any, None), + ) + + class PipelineProvider: + def create_pipeline_config( + self, *, preset_id: str, options: object + ) -> StreamInferencePipelineConfig: + del preset_id, options + return pipeline_config + + monkeypatch.setattr(application, "load_preset_catalog", lambda _: catalog) + monkeypatch.setattr( + application, "load_pipeline_provider", lambda _: PipelineProvider() + ) + runtime = application.create_runtime( + ApplicationArguments( + mode="mp4", + parser=argparse.ArgumentParser(), + argv=("--prompt", "A waterfall"), + ) + ) + + class Mode: + name = "test" + + def run(self, runtime: Runtime, drive_session: DriveSession) -> tuple[()]: + del runtime, drive_session + return () + + mode = Mode() + assert isinstance(mode, IOHandler) + runtime.initialize(device="cpu", io_handler=mode) + assert runtime.peek_steady_output_num_frames() == 3 + session = runtime.create_session(InferenceInput()) + assert isinstance(session, T2VSession) + result = session.step(InferenceInput()) + assert result.step_index == 0 + assert result.frame_count == 3 + assert result.metrics["step_ms"] == 1.0 + session.destroy() + runtime.destroy() + + assert calls == [ + "pipeline.init", + "pipeline.to", + "pipeline.eval", + "session.initialize_cache", + "session.pipeline.generate", + "session.pipeline.finalize", + "cache.close", + "pipeline.close", + ] + + +def test_total_blocks_is_only_a_finite_mode_default() -> None: + defaults = RuntimePresetOptions( + prompt="default prompt", + pixel_height=64, + pixel_width=96, + fps=12, + output_layout="tchw", + ) + + assert application._total_steps({"total_blocks": None}, defaults) is None + assert application._total_steps({"total_blocks": 4}, defaults) == 4 diff --git a/apps/t2v_app/tests/test_pipeline_provider.py b/apps/t2v_app/tests/test_pipeline_provider.py new file mode 100644 index 000000000..3cc1111f1 --- /dev/null +++ b/apps/t2v_app/tests/test_pipeline_provider.py @@ -0,0 +1,105 @@ +# 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. + +"""CPU tests for YAML preset and pipeline-provider resolution.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from t2v_app import application +from t2v_app.presets import load_preset_catalog + +from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig +from flashdreams.recipes.wan import Wan21TransformerConfig + +pytestmark = pytest.mark.ci_cpu + + +def test_packaged_yaml_constructs_default_pipeline_config() -> None: + catalog = load_preset_catalog() + preset_id, preset = catalog.resolve(None) + + config = application._create_pipeline_config(preset_id, preset) + + assert preset_id == "causal-forcing-wan2.1-t2v-1.3b-chunkwise" + assert config.name == preset_id + transformer = config.diffusion_model.transformer + scheduler = config.diffusion_model.scheduler + assert isinstance(transformer, Wan21TransformerConfig) + assert transformer.len_t == 3 + assert transformer.batch_shape == () + assert transformer.state_dict_transform is unwrap_generator_state_dict + assert isinstance(scheduler, FlowMatchSchedulerConfig) + assert scheduler.denoising_timesteps == [1000, 750, 500, 250] + assert preset.runtime.pixel_width == 832 + + +def test_catalog_rejects_incomplete_runtime_options(tmp_path: Path) -> None: + catalog_path = tmp_path / "presets.yaml" + catalog_path.write_text( + """ +schema_version: 1 +default_preset_id: test +presets: + test: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: test + total_blocks: 1 + pixel_height: 64 + pixel_width: 64 + fps: 16 + pipeline: {} +""".strip(), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="output_layout"): + load_preset_catalog(catalog_path) + + +def test_catalog_allows_total_blocks_to_be_omitted(tmp_path: Path) -> None: + catalog_path = tmp_path / "presets.yaml" + catalog_path.write_text( + """ +schema_version: 1 +default_preset_id: test +presets: + test: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + prompt: test + pixel_height: 64 + pixel_width: 64 + fps: 16 + output_layout: tchw + pipeline: {} +""".strip(), + encoding="utf-8", + ) + + _, preset = load_preset_catalog(catalog_path).resolve(None) + + assert preset.runtime.total_blocks is None + + +def test_catalog_reports_yaml_presets_for_unknown_id() -> None: + catalog = load_preset_catalog() + + with pytest.raises(ValueError, match="YAML presets"): + catalog.resolve("not-a-preset") diff --git a/apps/t2v_app/tests/test_webrtc.py b/apps/t2v_app/tests/test_webrtc.py new file mode 100644 index 000000000..45d9f73fb --- /dev/null +++ b/apps/t2v_app/tests/test_webrtc.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for T2V-specific WebRTC controls and recording.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, cast + +import pytest +import torch +from aiohttp import web + +from flashdreams.runtime import InferenceInput, OutputArtifact +from flashdreams.runtime.demo import DemoSpec, PreparedScenario, WebRTCOutputSpec +from flashdreams_runner import AppConfig, Runtime +from flashdreams_runner.webrtc import WebRTCMode +from t2v_app import runtime as runtime_module +from t2v_app import session as session_module +from t2v_app.runtime import T2VRuntime +from t2v_app.session import T2VScenario, T2VSession, T2VSessionDefaults +from t2v_app.webrtc import T2VWebRTCCustomization, T2VWebRTCSessionManager + +pytestmark = pytest.mark.ci_cpu + + +class _WebRuntime: + def __init__(self) -> None: + self.config = AppConfig( + model_id="t2v-app", + fps=12, + output_layout="tchw", + video_width=96, + video_height=64, + default_steps=2, + ) + self.latest_artifact = None + + def prepare_session_input( + self, + *, + prompt: str | None = None, + total_blocks: int | None = None, + ) -> InferenceInput: + return InferenceInput( + global_conditioning={ + "prompt": prompt or "default prompt", + "total_blocks": total_blocks, + "pixel_height": 64, + "pixel_width": 96, + "fps": 12, + } + ) + + def blocks_for_duration(self, duration_s: float) -> int: + return int(duration_s * 2) + + +def test_t2v_runtime_customizes_runner_webrtc_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Pipeline: + def to(self, device: object) -> "Pipeline": + assert device == "cpu" + return self + + def eval(self) -> "Pipeline": + return self + + class PipelineConfig: + def setup(self) -> Pipeline: + return Pipeline() + + monkeypatch.setattr(runtime_module, "StreamInferencePipeline", Pipeline) + runtime = T2VRuntime( + pipeline_config=cast(Any, PipelineConfig()), + session_defaults=T2VSessionDefaults( + prompt="default prompt", + total_blocks=2, + pixel_height=64, + pixel_width=96, + fps=12, + ), + config=AppConfig( + model_id="t2v-app", + fps=12, + output_layout="tchw", + video_width=96, + video_height=64, + default_steps=2, + ), + ) + mode = WebRTCMode(host="127.0.0.1", port=8080, device="cpu", world_rank=0) + + runtime.initialize(device="cpu", io_handler=mode) + + assert isinstance(mode._customization, T2VWebRTCCustomization) + runtime.destroy() + + +def test_t2v_customization_updates_prompt_duration_and_routes() -> None: + runtime = _WebRuntime() + customization = T2VWebRTCCustomization(runtime=cast(Any, runtime)) + initial_input = customization.prepare_initial_input() + assert initial_input.global_conditioning["total_blocks"] == 2 + + output = WebRTCOutputSpec( + host="127.0.0.1", + port=8080, + fps=12, + video_width=96, + video_height=64, + ) + spec = DemoSpec(model_id="t2v-app", input_mode="webrtc", output=output) + scenario = PreparedScenario(initial_inputs=initial_input) + + def provider_factory(spec: DemoSpec, scenario: PreparedScenario) -> Any: + del spec, scenario + return object() + + manager = customization.create_session_manager( + runtime=cast(Runtime, cast(object, runtime)), + output=output, + spec=spec, + scenario=scenario, + input_provider_factory=provider_factory, + ) + + assert isinstance(manager, T2VWebRTCSessionManager) + assert manager.is_runtime_ready() + assert manager._keep_connection_after_completed + manager.update_generation(prompt=" A waterfall ", duration_s=3.0) + prepared = manager._shared_scenario + assert prepared is not None + assert prepared.initial_inputs.global_conditioning["prompt"] == "A waterfall" + assert prepared.initial_inputs.global_conditioning["total_blocks"] == 6 + + resources = customization.create_app_resources(session_manager=manager) + assert resources.model_web_resource is not None + assert resources.model_web_resource.joinpath("adapter.js").is_file() + assert resources.configure_app is not None + app = web.Application() + resources.configure_app(app) + routes = { + resource.canonical + for route in app.router.routes() + if (resource := route.resource) is not None + } + assert { + "/api/t2v/config", + "/api/t2v/prompt", + "/api/t2v/download", + "/api/t2v/playback", + }.issubset(routes) + + +def test_t2v_session_records_completed_finite_generation( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + calls: list[str] = [] + + class Decoder: + spatial_compression_ratio = 8 + + class Cache: + def close(self) -> None: + calls.append("cache.close") + + class Pipeline: + decoder = Decoder() + + def initialize_cache(self, **kwargs: object) -> Cache: + del kwargs + return Cache() + + def get_num_output_frames(self, index: int) -> int: + del index + return 3 + + def generate(self, *, autoregressive_index: int, cache: object) -> torch.Tensor: + del autoregressive_index, cache + return torch.zeros((3, 3, 2, 2)) + + def finalize( + self, *, autoregressive_index: int, cache: object + ) -> dict[str, float]: + del autoregressive_index, cache + return {} + + class OutputTarget: + def __init__(self, *, output_path: Path, **kwargs: object) -> None: + del kwargs + self.output_path = output_path + + def open(self) -> None: + calls.append("output.open") + + def write(self, result: object) -> None: + del result + calls.append("output.write") + + def close(self) -> tuple[OutputArtifact, ...]: + calls.append("output.close") + return (OutputArtifact(kind="video/mp4", uri=str(self.output_path)),) + + monkeypatch.setattr(session_module, "StreamingVideoDecoder", Decoder) + monkeypatch.setattr(session_module, "Mp4VideoOutputTarget", OutputTarget) + recorded: list[tuple[Path, T2VScenario]] = [] + session = T2VSession( + pipeline=cast(Any, Pipeline()), + defaults=T2VSessionDefaults( + prompt="A waterfall", + total_blocks=1, + pixel_height=64, + pixel_width=96, + fps=12, + ), + initial_input=InferenceInput(), + output_layout="tchw", + record_artifact=lambda path, scenario: recorded.append((path, scenario)), + recording_directory=tmp_path, + ) + + assert session.next_step_request() is not None + session.step(InferenceInput()) + assert session.next_step_request() is None + session.destroy() + + assert calls == ["output.open", "output.write", "output.close", "cache.close"] + assert len(recorded) == 1 + assert recorded[0][0].parent == tmp_path + assert recorded[0][1].prompt == "A waterfall" diff --git a/flashdreams/flashdreams/core/checkpoint/remap.py b/flashdreams/flashdreams/core/checkpoint/remap.py index cce09795a..9d6610dd3 100644 --- a/flashdreams/flashdreams/core/checkpoint/remap.py +++ b/flashdreams/flashdreams/core/checkpoint/remap.py @@ -13,13 +13,50 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regex-based renaming of checkpoint state-dict keys.""" +"""Checkpoint state-dict unwrapping and key remapping.""" + +from __future__ import annotations import re +from collections.abc import Mapping +from typing import Any, cast from torch import Tensor +def unwrap_generator_state_dict(state_dict: dict[str, Any]) -> dict[str, Tensor]: + """Unwrap a generator checkpoint and strip root training prefixes. + + ``generator_ema`` takes precedence over ``generator`` when both containers + are present. A flat state dict passes through without envelope unwrapping. + + Args: + state_dict: Flat state dict or training checkpoint envelope. + + Returns: + Flat state dict without one ``model.`` or ``net.`` prefix followed by + an optional ``_fsdp_wrapped_module.`` prefix. + """ + if "generator_ema" in state_dict: + source = state_dict["generator_ema"] + elif "generator" in state_dict: + source = state_dict["generator"] + else: + source = state_dict + source = cast(Mapping[str, Tensor], source) + + transformed: dict[str, Tensor] = {} + for key, value in source.items(): + if key.startswith("model."): + key = key[len("model.") :] + elif key.startswith("net."): + key = key[len("net.") :] + if key.startswith("_fsdp_wrapped_module."): + key = key[len("_fsdp_wrapped_module.") :] + transformed[key] = value + return transformed + + def remap_checkpoint_keys( state_dict: dict[str, Tensor], mapping: dict[str, str] ) -> dict[str, Tensor]: diff --git a/flashdreams/flashdreams/core/pipeline_presets.py b/flashdreams/flashdreams/core/pipeline_presets.py new file mode 100644 index 000000000..2f37a36aa --- /dev/null +++ b/flashdreams/flashdreams/core/pipeline_presets.py @@ -0,0 +1,357 @@ +# 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. + +"""YAML pipeline-preset catalogs and declarative object-graph materialization.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from importlib import import_module +from pathlib import Path +from types import MappingProxyType +from typing import Any, Callable, Generic, Protocol, TypeVar, cast, runtime_checkable + +import yaml + +SCHEMA_VERSION = 1 +"""Supported YAML pipeline-preset schema version.""" + +_ROOT_FIELDS = {"schema_version", "default_preset_id", "presets"} +_PRESET_FIELDS = {"provider", "runtime", "pipeline"} + +RuntimeOptionsT = TypeVar("RuntimeOptionsT") +RuntimeOptionsT_co = TypeVar("RuntimeOptionsT_co", covariant=True) + + +class RuntimeOptionsParser(Protocol[RuntimeOptionsT_co]): + """Parse application-specific runtime options from one preset.""" + + def __call__(self, value: object, *, path: str) -> RuntimeOptionsT_co: + """Parse runtime options with a source path for validation errors.""" + ... + + +@dataclass(frozen=True, slots=True) +class PipelinePreset(Generic[RuntimeOptionsT]): + """Pipeline-provider selection and options for one preset.""" + + provider: str + """``module:attribute`` reference to a pipeline provider.""" + + runtime: RuntimeOptionsT + """Application-specific rollout and presentation defaults.""" + + pipeline: Mapping[str, object] + """Provider-owned pipeline construction options.""" + + +@dataclass(frozen=True, slots=True) +class PresetCatalog(Generic[RuntimeOptionsT]): + """Validated collection of named pipeline presets.""" + + default_preset_id: str + """Preset selected when a caller omits the preset identity.""" + + presets: Mapping[str, PipelinePreset[RuntimeOptionsT]] + """Preset definitions keyed by stable pipeline identity.""" + + def resolve( + self, preset_id: str | None + ) -> tuple[str, PipelinePreset[RuntimeOptionsT]]: + """Resolve an optional preset identity against this catalog.""" + selected = self.default_preset_id if preset_id is None else preset_id + try: + return selected, self.presets[selected] + except KeyError as exc: + available = ", ".join(sorted(self.presets)) + raise ValueError( + f"Unknown pipeline preset {selected!r}. YAML presets: {available}." + ) from exc + + +@runtime_checkable +class PipelineProvider(Protocol): + """Construct a pipeline config from preset-owned options.""" + + def create_pipeline_config( + self, + *, + preset_id: str, + options: Mapping[str, object], + ) -> object: + """Create a pipeline config for one resolved preset. + + Args: + preset_id: Selected preset identity. + options: Provider-owned options loaded from the preset YAML. + + Returns: + Pipeline config for application-specific validation. + """ + ... + + +class ObjectGraphPipelineProvider: + """Materialize nested Python configs from a YAML object graph. + + A mapping containing ``_target`` imports and calls the named object with + the remaining mapping entries as keyword arguments. ``_ref`` imports an + object without calling it, and ``_tuple`` preserves tuple-valued config + fields that YAML otherwise represents as lists. + """ + + def create_pipeline_config( + self, + *, + preset_id: str, + options: Mapping[str, object], + ) -> object: + """Materialize one pipeline config object graph.""" + return materialize_object_graph( + options, + path=f"presets.{preset_id}.pipeline", + ) + + +def load_pipeline_preset_catalog( + path: str | Path, + *, + runtime_options_parser: RuntimeOptionsParser[RuntimeOptionsT], +) -> PresetCatalog[RuntimeOptionsT]: + """Load and validate a pipeline-preset catalog from a YAML file. + + Args: + path: YAML catalog path. + runtime_options_parser: Application-specific runtime-options parser. + + Returns: + Validated preset catalog. + """ + source = Path(path) + return parse_pipeline_preset_catalog( + source.read_text(encoding="utf-8"), + source_name=str(source), + runtime_options_parser=runtime_options_parser, + ) + + +def parse_pipeline_preset_catalog( + raw_text: str, + *, + source_name: str, + runtime_options_parser: RuntimeOptionsParser[RuntimeOptionsT], +) -> PresetCatalog[RuntimeOptionsT]: + """Parse and validate a pipeline-preset catalog from YAML text. + + Args: + raw_text: YAML catalog contents. + source_name: Source label included in validation errors. + runtime_options_parser: Application-specific runtime-options parser. + + Returns: + Validated preset catalog. + """ + raw = yaml.safe_load(raw_text) + root = _mapping(raw, path=source_name) + _require_exact_fields(root, expected=_ROOT_FIELDS, path=source_name) + if root["schema_version"] != SCHEMA_VERSION: + raise ValueError( + f"{source_name}.schema_version must be {SCHEMA_VERSION}, " + f"got {root['schema_version']!r}." + ) + + default_preset_id = _nonempty_string( + root["default_preset_id"], path=f"{source_name}.default_preset_id" + ) + raw_presets = _mapping(root["presets"], path=f"{source_name}.presets") + if not raw_presets: + raise ValueError(f"{source_name}.presets must not be empty.") + + presets: dict[str, PipelinePreset[RuntimeOptionsT]] = {} + for raw_name, raw_preset in raw_presets.items(): + name = _nonempty_string(raw_name, path=f"{source_name}.presets key") + preset_path = f"{source_name}.presets.{name}" + preset = _mapping(raw_preset, path=preset_path) + _require_exact_fields(preset, expected=_PRESET_FIELDS, path=preset_path) + provider = _nonempty_string(preset["provider"], path=f"{preset_path}.provider") + runtime = runtime_options_parser( + preset["runtime"], + path=f"{preset_path}.runtime", + ) + pipeline = _mapping(preset["pipeline"], path=f"{preset_path}.pipeline") + presets[name] = PipelinePreset( + provider=provider, + runtime=runtime, + pipeline=MappingProxyType(dict(pipeline)), + ) + + if default_preset_id not in presets: + raise ValueError( + f"{source_name}.default_preset_id {default_preset_id!r} is not " + "defined under presets." + ) + return PresetCatalog( + default_preset_id=default_preset_id, + presets=MappingProxyType(presets), + ) + + +def load_pipeline_provider(reference: str) -> PipelineProvider: + """Load a pipeline-provider instance from ``module:attribute``.""" + candidate = resolve_reference(reference) + provider = candidate() if isinstance(candidate, type) else candidate + if not isinstance(provider, PipelineProvider): + raise TypeError( + f"Pipeline provider {reference!r} resolved to " + f"{type(provider).__name__}, which does not implement " + "create_pipeline_config()." + ) + return provider + + +def resolve_reference(reference: str) -> object: + """Resolve one ``module:attribute`` reference.""" + module_name, separator, attribute_path = reference.partition(":") + if not separator or not module_name or not attribute_path: + raise ValueError( + f"Invalid Python reference {reference!r}; expected module:attribute." + ) + value: object = import_module(module_name) + for attribute in attribute_path.split("."): + try: + value = getattr(value, attribute) + except AttributeError as exc: + raise ValueError( + f"Python reference {reference!r} has no attribute {attribute!r}." + ) from exc + return value + + +def materialize_object_graph(value: object, *, path: str = "root") -> object: + """Materialize declarative ``_target``, ``_ref``, and ``_tuple`` nodes. + + Args: + value: YAML-decoded object graph. + path: Source path included in validation errors. + + Returns: + Recursively materialized Python object. + """ + if isinstance(value, list): + return [ + materialize_object_graph(item, path=f"{path}[{index}]") + for index, item in enumerate(value) + ] + if not isinstance(value, Mapping): + return value + + mapping = cast(dict[str, object], dict(value)) + if any(not isinstance(key, str) for key in mapping): + raise TypeError(f"{path} keys must be strings.") + reserved = {key for key in mapping if key.startswith("_")} + if "_ref" in mapping: + if reserved != {"_ref"} or len(mapping) != 1: + raise ValueError(f"{path}._ref cannot be combined with other keys.") + reference = mapping["_ref"] + if not isinstance(reference, str): + raise TypeError(f"{path}._ref must be a string.") + return resolve_reference(reference) + + if "_tuple" in mapping: + if reserved != {"_tuple"} or len(mapping) != 1: + raise ValueError(f"{path}._tuple cannot be combined with other keys.") + items = mapping["_tuple"] + if not isinstance(items, list): + raise TypeError(f"{path}._tuple must contain a YAML list.") + return tuple( + materialize_object_graph(item, path=f"{path}._tuple[{index}]") + for index, item in enumerate(items) + ) + + if "_target" not in mapping: + if reserved: + names = ", ".join(sorted(reserved)) + raise ValueError(f"Unsupported reserved keys at {path}: {names}.") + return { + str(key): materialize_object_graph(item, path=f"{path}.{key}") + for key, item in mapping.items() + } + + if reserved != {"_target"}: + names = ", ".join(sorted(reserved - {"_target"})) + raise ValueError(f"Unsupported reserved keys at {path}: {names}.") + reference = mapping["_target"] + if not isinstance(reference, str): + raise TypeError(f"{path}._target must be a string.") + target = resolve_reference(reference) + if not callable(target): + raise TypeError(f"{path}._target {reference!r} is not callable.") + target_callable = cast(Callable[..., object], target) + kwargs = { + str(key): materialize_object_graph(item, path=f"{path}.{key}") + for key, item in mapping.items() + if key != "_target" + } + try: + return target_callable(**kwargs) + except Exception as exc: + raise ValueError( + f"Failed to construct {path} with {reference!r}: {exc}" + ) from exc + + +def _mapping(value: object, *, path: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise TypeError(f"{path} must be a mapping, got {type(value).__name__}.") + mapping = cast(dict[object, object], dict(cast(Any, value))) + if any(not isinstance(key, str) for key in mapping): + raise TypeError(f"{path} keys must be strings.") + return cast(dict[str, object], mapping) + + +def _require_exact_fields( + value: Mapping[str, object], *, expected: set[str], path: str +) -> None: + fields = {str(key) for key in value} + missing = expected - fields + unknown = fields - expected + if missing or unknown: + details: list[str] = [] + if missing: + details.append(f"missing {sorted(missing)}") + if unknown: + details.append(f"unknown {sorted(unknown)}") + raise ValueError(f"Invalid fields at {path}: {'; '.join(details)}.") + + +def _nonempty_string(value: object, *, path: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise TypeError(f"{path} must be a non-empty string.") + return value.strip() + + +__all__ = [ + "ObjectGraphPipelineProvider", + "PipelinePreset", + "PipelineProvider", + "PresetCatalog", + "RuntimeOptionsParser", + "load_pipeline_preset_catalog", + "load_pipeline_provider", + "materialize_object_graph", + "parse_pipeline_preset_catalog", + "resolve_reference", +] diff --git a/flashdreams/flashdreams/runtime/demo/drivers.py b/flashdreams/flashdreams/runtime/demo/drivers.py index 27179ac32..dfbe4dd41 100644 --- a/flashdreams/flashdreams/runtime/demo/drivers.py +++ b/flashdreams/flashdreams/runtime/demo/drivers.py @@ -7,6 +7,7 @@ import asyncio import inspect +from collections.abc import Callable from typing import Any, cast from flashdreams.runtime.interfaces import InferenceSession @@ -446,8 +447,15 @@ async def run_demo_session_async( run_mode: RunMode, pipeline: StepPipeline, reservation: SessionReservation | None = None, + model_input_provider_factory: ( + Callable[[DemoSpec, PreparedScenario], ModelInputProvider] | None + ) = None, ) -> RunResult: - """Run one prepared async/realtime demo session through a selected run mode.""" + """Run one prepared async/realtime demo session through a selected run mode. + + ``model_input_provider_factory`` overrides the adapter's provider factory + for callers that already own a runtime and prepared scenario. + """ if reservation is None: reservation = context.admission.try_reserve() if reservation is None: @@ -459,7 +467,9 @@ async def run_demo_session_async( session_edges: SessionEdges | None = None try: try: - create_provider = getattr(adapter, "create_model_input_provider") + create_provider = model_input_provider_factory + if create_provider is None: + create_provider = getattr(adapter, "create_model_input_provider") provider = await context.host.call_async(create_provider, spec, scenario) run_mode.validate_session( spec=spec, diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index 149f83d4f..b0785be07 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -680,9 +680,13 @@ def __init__( shared_spec: DemoSpec | None = None, shared_spec_factory: Callable[[Any], DemoSpec] | None = None, shared_scenario: PreparedScenario | None = None, + shared_model_input_provider_factory: ( + Callable[[DemoSpec, PreparedScenario], ModelInputProvider] | None + ) = None, shared_pipeline_factory: Callable[[], StepPipeline] | None = None, legacy_segment_resampler_factory: Callable[..., Any] | None = None, keep_connection_after_completed: bool = False, + runtime_ready: bool = False, ) -> None: if client_liveness_timeout_s <= 0: raise ValueError("client_liveness_timeout_s must be > 0") @@ -699,7 +703,9 @@ def __init__( self.fatal_generation_errors = fatal_generation_errors self.client_liveness_timeout_s = client_liveness_timeout_s self._runtime = runtime - self._runtime_ready = False + # Runner-hosted applications may finish model initialization before + # the WebRTC server lifecycle begins. + self._runtime_ready = runtime_ready self._warmup_complete = False self._active_session: ManagedWebRTCSession | None = None self._preload_lock = asyncio.Lock() @@ -713,6 +719,7 @@ def __init__( self._shared_spec = shared_spec self._shared_spec_factory = shared_spec_factory self._shared_scenario = shared_scenario + self._shared_model_input_provider_factory = shared_model_input_provider_factory self._shared_pipeline_factory = shared_pipeline_factory self._keep_connection_after_completed = keep_connection_after_completed self._shared_video_encoder: VideoEncoder | None = None @@ -760,9 +767,16 @@ def _make_legacy_segment_resampler_at_fps( kwargs["supported_keys"] = supported_control_keys return factory(**kwargs) + def _uses_shared_demo_path(self) -> bool: + return ( + self._shared_adapter is not None + or self._shared_model_input_provider_factory is not None + ) + def _needs_legacy_segment_metadata(self) -> bool: - return self._shared_adapter is None and not _runtime_drives_inference_session( - self._runtime + return ( + not self._uses_shared_demo_path() + and not _runtime_drives_inference_session(self._runtime) ) def _effective_supported_control_keys(self) -> frozenset[str] | None: @@ -894,7 +908,7 @@ async def _reset_runtime_for_session( ) -> None: reset = getattr(context.host.runtime, "reset_for_new_session", None) if not callable(reset): - if self._shared_adapter is not None: + if self._uses_shared_demo_path(): return raise RuntimeError("WebRTC runtime adapter cannot reset sessions.") await context.host.call_async(reset, session_input) @@ -1342,7 +1356,7 @@ async def preload_runtime(self) -> None: elif self._shared_host is not None: await asyncio.to_thread(self._shared_host.preload) self._runtime_ready = True - self._initialize_shared_video_encoder() + self._initialize_shared_video_encoder() if not self._warmup_complete: await self._run_loopback_warmup_session( num_chunks=self.runtime_config.warmup_chunks @@ -1886,11 +1900,18 @@ async def _run_realtime_driver_session( ) try: while not managed_session.closed: - adapter = self._shared_adapter + adapter: Any = self._shared_adapter spec = self._shared_spec scenario = self._shared_scenario spec_factory = self._shared_spec_factory - if adapter is None or spec is None: + provider_factory = self._shared_model_input_provider_factory + if provider_factory is not None: + if spec is None: + raise RuntimeError( + "Direct WebRTC input providers require shared_spec." + ) + adapter = self._runtime + elif adapter is None or spec is None: adapter = _LegacyWebRTCDemoAdapter( runtime=self._runtime, identity=self.identity, @@ -1901,12 +1922,16 @@ async def _run_realtime_driver_session( spec = spec_factory(session_input) scenario = None if scenario is None: - scenario = adapter.prepare_scenario(spec) + if provider_factory is not None: + raise RuntimeError( + "Direct WebRTC input providers require shared_scenario." + ) + scenario = cast(Any, adapter).prepare_scenario(spec) result = await run_demo_session_async( context=context, spec=spec, scenario=scenario, - adapter=adapter, + adapter=cast(Any, adapter), run_mode=run_mode, pipeline=( self._shared_pipeline_factory() @@ -1914,6 +1939,7 @@ async def _run_realtime_driver_session( else StepPipeline() ), reservation=managed_session.reservation, + model_input_provider_factory=provider_factory, ) managed_session.reservation = None if result.status != "completed": diff --git a/flashdreams/tests/test_checkpoint_loading.py b/flashdreams/tests/test_checkpoint_loading.py index f3b129bd1..2e4303df2 100644 --- a/flashdreams/tests/test_checkpoint_loading.py +++ b/flashdreams/tests/test_checkpoint_loading.py @@ -14,9 +14,55 @@ import torch from safetensors.torch import save_file as save_safetensors_file +from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict + pytestmark = pytest.mark.ci_cpu +@pytest.mark.parametrize("container", ["generator_ema", "generator"]) +def test_unwrap_generator_state_dict_strips_training_prefixes( + container: str, +) -> None: + """Unwrap generator containers and strip their root training prefixes.""" + model_weight = torch.tensor(1.0) + net_bias = torch.tensor(2.0) + fsdp_scale = torch.tensor(3.0) + untouched = torch.tensor(4.0) + + actual = unwrap_generator_state_dict( + { + container: { + "model.weight": model_weight, + "net.bias": net_bias, + "_fsdp_wrapped_module.scale": fsdp_scale, + "untouched": untouched, + } + } + ) + + assert actual == { + "weight": model_weight, + "bias": net_bias, + "scale": fsdp_scale, + "untouched": untouched, + } + + +def test_unwrap_generator_state_dict_prefers_ema_container() -> None: + """Prefer EMA parameters when both generator containers are present.""" + generator = torch.tensor(1.0) + generator_ema = torch.tensor(2.0) + + actual = unwrap_generator_state_dict( + { + "generator": {"model.weight": generator}, + "generator_ema": {"model._fsdp_wrapped_module.weight": generator_ema}, + } + ) + + assert actual == {"weight": generator_ema} + + def test_local_safetensors_uses_file_backed_loader( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/flashdreams/tests/test_demo_runtime_run_modes.py b/flashdreams/tests/test_demo_runtime_run_modes.py index 8742eb0e9..1daff8839 100644 --- a/flashdreams/tests/test_demo_runtime_run_modes.py +++ b/flashdreams/tests/test_demo_runtime_run_modes.py @@ -38,6 +38,7 @@ NullOutputSpec, OutputDecision, PreparedScenario, + PreparedStep, ProviderCapabilities, RunContext, RunModeCapabilities, @@ -165,6 +166,51 @@ async def test_fake_webrtc_offer_reserves_before_prepare_or_negotiation() -> Non assert mode.created_edges[0].is_closed +@pytest.mark.asyncio +async def test_async_session_uses_explicit_model_input_provider_factory() -> None: + spec = DemoSpec( + model_id="fake-demo", + input_mode="keyboard-driving", + output=WebRTCOutputSpec(port=8081), + ) + adapter = _FakeAdapter() + provider = _FakeProvider() + mode = _FakeRunMode(name="webrtc", driver=_ClosingAsyncDriver()) + context = mode.create_run_context( + spec=spec, + adapter=adapter, + host=RuntimeHost(_UnusedRuntime()), + model_warmup_plan=ModelWarmupPlan(), + ) + scenario = adapter.prepare_scenario(spec) + factory_calls: list[tuple[DemoSpec, PreparedScenario]] = [] + + def create_provider( + factory_spec: DemoSpec, + factory_scenario: PreparedScenario, + ) -> _FakeProvider: + factory_calls.append((factory_spec, factory_scenario)) + return provider + + try: + result = await run_demo_session_async( + context=context, + spec=spec, + scenario=scenario, + adapter=adapter, + run_mode=mode, + pipeline=StepPipeline(), + model_input_provider_factory=create_provider, + ) + finally: + context.host.close() + + assert result.status == "completed" + assert factory_calls == [(spec, scenario)] + assert adapter.providers == [] + assert provider.close_count == 1 + + @pytest.mark.asyncio async def test_async_session_cancellation_shields_pre_edge_provider_cleanup() -> None: spec = DemoSpec( @@ -499,6 +545,21 @@ class _FakeProvider: def __init__(self) -> None: self.close_count = 0 + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput() + + 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: + del inputs + def close(self) -> None: self.close_count += 1 diff --git a/flashdreams/tests/test_pipeline_presets.py b/flashdreams/tests/test_pipeline_presets.py new file mode 100644 index 000000000..5db258ddc --- /dev/null +++ b/flashdreams/tests/test_pipeline_presets.py @@ -0,0 +1,117 @@ +# 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. + +"""CPU tests for shared YAML pipeline-preset parsing and materialization.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import pytest + +from flashdreams.core.pipeline_presets import ( + ObjectGraphPipelineProvider, + load_pipeline_preset_catalog, + load_pipeline_provider, + materialize_object_graph, +) + +pytestmark = pytest.mark.ci_cpu + + +def _parse_runtime_options(value: object, *, path: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise TypeError(f"{path} must be a mapping.") + return dict(cast(Mapping[str, object], value)) + + +def test_load_pipeline_preset_catalog_parses_generic_schema(tmp_path: Path) -> None: + path = tmp_path / "presets.yaml" + path.write_text( + """ +schema_version: 1 +default_preset_id: example +presets: + example: + provider: flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider + runtime: + fps: 12 + pipeline: + enabled: true +""".strip(), + encoding="utf-8", + ) + + catalog = load_pipeline_preset_catalog( + path, + runtime_options_parser=_parse_runtime_options, + ) + preset_id, preset = catalog.resolve(None) + + assert preset_id == "example" + assert preset.runtime == {"fps": 12} + assert preset.pipeline == {"enabled": True} + + +def test_load_pipeline_preset_catalog_rejects_unknown_root_field( + tmp_path: Path, +) -> None: + path = tmp_path / "presets.yaml" + path.write_text( + """ +schema_version: 1 +default_preset_id: example +presets: {} +unexpected: true +""".strip(), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="unknown.*unexpected"): + load_pipeline_preset_catalog( + path, + runtime_options_parser=_parse_runtime_options, + ) + + +def test_object_graph_provider_materializes_shared_declarative_nodes() -> None: + provider = load_pipeline_provider( + "flashdreams.core.pipeline_presets:ObjectGraphPipelineProvider" + ) + + config = provider.create_pipeline_config( + preset_id="example", + options={ + "_target": "types:SimpleNamespace", + "callback": {"_ref": "builtins:len"}, + "shape": {"_tuple": [1, 2, 3]}, + }, + ) + + assert isinstance(provider, ObjectGraphPipelineProvider) + assert isinstance(config, SimpleNamespace) + assert config.callback is len + assert config.shape == (1, 2, 3) + + +def test_materialize_object_graph_rejects_mixed_reference_node() -> None: + with pytest.raises(ValueError, match="cannot be combined"): + materialize_object_graph( + {"_ref": "builtins:len", "other": True}, + path="preset.pipeline.callback", + ) diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index bafe2b1a8..edbed1d63 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -19,7 +19,12 @@ UserInputEvent, UserInputs, ) -from flashdreams.runtime.demo import RunResult +from flashdreams.runtime.demo import ( + DemoSpec, + PreparedScenario, + RunResult, + WebRTCOutputSpec, +) from flashdreams.runtime.keyboard import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc import manager as manager_module from flashdreams.serving.webrtc.encoders import ChunkDeliveryResult @@ -1108,6 +1113,66 @@ def peek_steady_output_num_frames(self) -> int: assert chunk_done[0]["model"] == "fake-model" +@pytest.mark.asyncio +async def test_realtime_driver_session_forwards_direct_provider_factory( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + async def fake_run_demo_session_async(**kwargs: Any) -> RunResult: + captured.update(kwargs) + return RunResult(status="completed") + + def provider_factory( + spec: DemoSpec, + scenario: PreparedScenario, + ) -> Any: + del spec, scenario + raise AssertionError("The fake session helper must not create the provider.") + + monkeypatch.setattr( + manager_module, + "run_demo_session_async", + fake_run_demo_session_async, + ) + runtime = SimpleNamespace(close=lambda: None) + spec = DemoSpec( + model_id="fake-model", + input_mode="webrtc", + output=WebRTCOutputSpec(), + ) + scenario = PreparedScenario(initial_inputs=InferenceInput()) + manager = _make_manager( + _BaseTestManager, + runtime, + shared_host=manager_module.RuntimeHost(cast(Any, runtime)), + shared_spec=spec, + shared_scenario=scenario, + shared_model_input_provider_factory=provider_factory, + ) + assert not manager._needs_legacy_segment_metadata() + context = manager._shared_run_context(asyncio.get_running_loop()) + await manager._reset_runtime_for_session(context=context, session_input=None) + reservation = context.admission.try_reserve() + assert reservation is not None + managed, _video_track, _peer, _channel = _managed_session(runtime) + managed.reservation = reservation + manager._active_session = managed + + await manager._run_realtime_driver_session( + managed_session=managed, + context=context, + session_input=None, + ) + + assert captured["adapter"] is runtime + assert captured["spec"] is spec + assert captured["scenario"] is scenario + assert captured["model_input_provider_factory"] is provider_factory + assert not manager.has_active_session() + context.host.close() + + @pytest.mark.asyncio async def test_realtime_driver_session_reports_non_completed_result( monkeypatch: pytest.MonkeyPatch, diff --git a/flashdreams_runner/README.md b/flashdreams_runner/README.md new file mode 100644 index 000000000..35f3d45ef --- /dev/null +++ b/flashdreams_runner/README.md @@ -0,0 +1,158 @@ +# FlashDreams Runner + +`flashdreams-runner` is the model-neutral shell for FlashDreams applications. +It loads an application from the active Python environment, creates its runtime, +initializes that runtime with a selected I/O mode, creates sessions, and drives +their generation loops. + +```bash +uv run flashdreams-runner t2v-app mp4 --output o.mp4 --prompt "A waterfall" +uv run flashdreams-runner t2v-app replay --output o.mp4 --prompt "A waterfall" +uv run flashdreams-runner t2v-app webrtc +uv run flashdreams-runner t2v-app none --steps 4 --prompt "A waterfall" +``` + +## Architecture + +```text +application module flashdreams-runner I/O mode ++------------------+ +--------------------------+ +---------------+ +| create_runtime() | ----> | initialize Runtime | ----> | Replay / MP4 | +| | | create Session | | WebRTC | +| Runtime | | | | None | +| model weights | | loop: | +-------+-------+ +| global state | | input = mode.read() | | +| | | output = Session. |<--------------+ +| Session |<------| step(input) | +| prompt/cache | | mode.write(output) |-------------->+ +| game state | | destroy Session/Runtime | ++------------------+ +--------------------------+ +``` + +The application owns inference. The runner owns orchestration and I/O. + +## Application ABI + +An installed application module exposes one function: + +```python +create_runtime(arguments: flashdreams_runner.ApplicationArguments) \ + -> flashdreams_runner.Runtime +``` + +The function adds application-specific arguments to `arguments.parser`, calls +`arguments.parse_args()`, resolves application configuration, and returns an +uninitialized runtime. This single factory is the ABI between an application +package and `flashdreams-runner`. + +`Runtime` owns model weights and other one-time or process-global state: + +- `config` exposes the application's `AppConfig` to runner-owned modes. +- `initialize(device=..., io_handler=...)` performs model construction and + one-time setup. +- `create_session(initial_input)` creates isolated per-user state. +- `destroy()` releases model and process resources. + +`Session` owns the application loop implementation and all per-user state, such +as prompts, K/V caches, world state, and step counters: + +- `step(inputs)` performs exactly one application iteration and returns a + `StepResult`. +- `destroy()` releases session-local resources. + +The base classes also expose compatibility spellings for the shared +FlashDreams serving stack, so WebRTC consumes an application runtime directly +without a runner-specific adapter. + +An application that needs more than the generic viewer can configure the +runner's `WebRTCMode` during `Runtime.initialize()`. The customization supplies +initial session input, a specialized session manager, packaged browser assets, +and application HTTP routes; the runner still owns WebRTC transport and server +lifecycle. + +## I/O modes + +Modes are runner-owned I/O handlers. They never construct model pipelines or +implement application generation logic. + +| Mode | Input/output behavior | +|---|---| +| `mp4` | Compatibility name for a finite replay written to MP4. | +| `replay` | Runs a finite deterministic input sequence and writes MP4. | +| `webrtc` | Creates a live server and lets applications add model-specific browser controls and routes. | +| `none` | Runs a finite input sequence and discards output. | + +`--steps` overrides the finite iteration count for `mp4`, `replay`, and `none`. +If omitted, those modes use `Runtime.config.default_steps`. Native-window and +additional interactive handlers can implement the same `IOHandler` boundary. + +## Minimal application + +```python +from flashdreams.runtime import InferenceInput, StepResult +from flashdreams_runner import ( + ApplicationArguments, + AppConfig, + IOHandler, + Runtime, + Session, +) + + +class MySession(Session): + def __init__(self, model, prompt: str) -> None: + self.model = model + self.prompt = prompt + self.cache = model.create_cache(prompt) + self._step_index = 0 + + @property + def step_index(self) -> int: + return self._step_index + + def step(self, inputs: InferenceInput) -> StepResult: + video = self.model.generate(self.cache, inputs.step) + result = StepResult.from_video_chunk( + step_index=self._step_index, + video_chunk=video, + layout="tchw", + ) + self._step_index += 1 + return result + + def destroy(self) -> None: + self.cache = None + + +class MyRuntime(Runtime): + def __init__(self, prompt: str) -> None: + self.prompt = prompt + self.model = None + + @property + def config(self) -> AppConfig: + return AppConfig( + model_id="my-app", + fps=24, + output_layout="tchw", + video_width=832, + video_height=480, + default_steps=4, + ) + + def initialize(self, *, device: str, io_handler: IOHandler) -> None: + del io_handler + self.model = load_model(device) + + def create_session(self, initial_input=None) -> Session: + return MySession(self.model, self.prompt) + + def destroy(self) -> None: + self.model = None + + +def create_runtime(arguments: ApplicationArguments) -> Runtime: + arguments.parser.add_argument("--prompt", required=True) + options = arguments.parse_args() + return MyRuntime(prompt=options.prompt) +``` diff --git a/flashdreams_runner/__init__.py b/flashdreams_runner/__init__.py new file mode 100644 index 000000000..52d40a9ce --- /dev/null +++ b/flashdreams_runner/__init__.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Public ABI for FlashDreams applications, runtimes, and sessions.""" + +from .contracts import ( + AppConfig, + Application, + ApplicationArguments, + DriveSession, + IOHandler, + InputHandler, + OutputHandler, + Runtime, + Session, +) + +__all__ = [ + "AppConfig", + "Application", + "ApplicationArguments", + "DriveSession", + "IOHandler", + "InputHandler", + "OutputHandler", + "Runtime", + "Session", +] diff --git a/flashdreams_runner/cli.py b/flashdreams_runner/cli.py new file mode 100644 index 000000000..626c80d93 --- /dev/null +++ b/flashdreams_runner/cli.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command-line shell for independently installed FlashDreams applications.""" + +from __future__ import annotations + +import argparse +import importlib +from contextlib import ExitStack +from dataclasses import dataclass +from importlib import metadata +from typing import Sequence + +import torch + +from flashdreams.runtime import OutputArtifact, StepResult +from flashdreams.runtime.demo.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) + +from .contracts import ( + Application, + ApplicationArguments, + InputHandler, + OutputHandler, + Runtime, + Session, +) +from .modes import MODE_NAMES, add_mode_arguments, create_io_handler + + +@dataclass(frozen=True, slots=True) +class _ApplicationAndMode: + """Top-level route parsed before loading an application.""" + + application: str + """Installed application distribution name.""" + + mode: str + """Selected runner I/O mode.""" + + remaining_argv: tuple[str, ...] + """Arguments delegated to the selected mode and application.""" + + +@dataclass(frozen=True, slots=True) +class _Environment: + """Initialized device and distributed process information.""" + + device: str + """Resolved runtime device for this process.""" + + world_rank: int + """Global distributed rank.""" + + world_size: int + """Number of distributed processes.""" + + +def build_parser(application: str, mode: str) -> argparse.ArgumentParser: + """Build a parser containing runner and selected-mode arguments. + + Args: + application: Application name displayed in command usage. + mode: Selected I/O mode. + + Returns: + Parser for the application factory to extend and invoke. + """ + parser = argparse.ArgumentParser(prog=f"flashdreams-runner {application} {mode}") + add_mode_arguments(parser, mode) + return parser + + +def load_application(distribution_name: str) -> Application: + """Load an installed module that satisfies the application ABI.""" + try: + distribution = metadata.distribution(distribution_name) + except metadata.PackageNotFoundError as exc: + raise ValueError( + f"Application distribution {distribution_name!r} is not installed." + ) from exc + + package_names = metadata.packages_distributions() + candidates = [ + name + for name, distributions in package_names.items() + if distribution.metadata["Name"] in distributions + ] + candidates.append(distribution_name.replace("-", "_")) + incompatible: list[str] = [] + for candidate in dict.fromkeys(candidates): + try: + module = importlib.import_module(candidate) + except ModuleNotFoundError as exc: + if exc.name != candidate: + raise + continue + if isinstance(module, Application): + return module + incompatible.append(candidate) + if incompatible: + names = ", ".join(repr(name) for name in incompatible) + raise TypeError( + f"Application distribution {distribution_name!r} exposes module(s) " + f"{names}, but none satisfy Application. An application must define " + "create_runtime(arguments)." + ) + raise ValueError( + f"Application {distribution_name!r} does not expose an importable module." + ) + + +def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: + """Create an application runtime and dispatch it to one I/O mode. + + Args: + argv: Command-line arguments; ``None`` reads the process arguments. + + Returns: + Persistent artifacts produced by the selected mode. + """ + route = _parse_application_and_mode(argv) + application = load_application(route.application) + arguments = ApplicationArguments( + mode=route.mode, + parser=build_parser(route.application, route.mode), + argv=route.remaining_argv, + ) + runtime = _require_runtime( + application.create_runtime(arguments), + application_name=route.application, + ) + options = arguments.options + environment = _initialize_environment(options.device) + io_handler = create_io_handler( + route.mode, + options, + device=environment.device, + world_rank=environment.world_rank, + ) + try: + runtime.initialize( + device=environment.device, + io_handler=io_handler, + ) + return io_handler.run(runtime, _drive_session) + finally: + runtime.destroy() + + +def _parse_application_and_mode( + argv: Sequence[str] | None, +) -> _ApplicationAndMode: + """Parse application and mode while preserving all remaining arguments.""" + parser = argparse.ArgumentParser(prog="flashdreams-runner", add_help=False) + parser.add_argument("application", help="Installed application distribution") + parser.add_argument("mode", choices=MODE_NAMES) + args, remaining_argv = parser.parse_known_args(argv) + return _ApplicationAndMode( + application=args.application, + mode=args.mode, + remaining_argv=tuple(remaining_argv), + ) + + +def _require_runtime(value: object, *, application_name: str) -> Runtime: + """Validate the application factory result before initialization.""" + if not isinstance(value, Runtime): + raise TypeError( + f"Application {application_name!r} create_runtime() returned " + f"{type(value).__name__}, expected Runtime." + ) + return value + + +def _drive_session( + runtime: Runtime, + input_handler: InputHandler, + output_handler: OutputHandler, +) -> tuple[OutputArtifact, ...]: + """Drive one application session through a pair of I/O handlers.""" + with ExitStack() as resources: + input_handler.open() + resources.callback(input_handler.close) + + output_handler.open(runtime.config) + output_closed = False + + def close_output() -> None: + if not output_closed: + output_handler.close() + + resources.callback(close_output) + session = runtime.create_session(input_handler.initial_input()) + if not isinstance(session, Session): + raise TypeError( + "Runtime.create_session() must return Session, got " + f"{type(session).__name__}." + ) + resources.callback(session.destroy) + + while (inputs := input_handler.read()) is not None: + result = session.step(inputs) + if not isinstance(result, StepResult): + raise TypeError( + "Session.step() must return StepResult, got " + f"{type(result).__name__}." + ) + output_handler.write(result) + + artifacts = tuple(output_handler.close()) + output_closed = True + return artifacts + + +def _initialize_environment(device: str) -> _Environment: + """Initialize logging, CUDA placement, and distributed process state.""" + if torch.device(device).type == "cuda": + context = initialize_cuda_distributed(default_device=device) + return _Environment( + device=str(context.device), + world_rank=context.world_rank, + world_size=context.world_size, + ) + configure_logging(world_rank=0) + return _Environment(device=str(torch.device(device)), world_rank=0, world_size=1) + + +def main() -> None: + """Run the console-script entry point.""" + artifacts = run() + # Persistent modes return artifact URIs; live and headless modes return none. + for artifact in artifacts: + print(artifact.uri) + + +__all__ = ["build_parser", "load_application", "main", "run"] diff --git a/flashdreams_runner/contracts.py b/flashdreams_runner/contracts.py new file mode 100644 index 000000000..fe6102e72 --- /dev/null +++ b/flashdreams_runner/contracts.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Application, runtime, session, and I/O mode contracts.""" + +from __future__ import annotations + +import argparse +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.runtime import InferenceInput, OutputArtifact, StepRequest, StepResult + +if TYPE_CHECKING: + from flashdreams.runtime.demo import SessionInfo + + +@dataclass(frozen=True, kw_only=True, slots=True) +class AppConfig: + """Application output configuration consumed by runner-owned I/O modes.""" + + model_id: str + """Stable application or model identity.""" + + fps: int | float + """Output video frame rate.""" + + output_layout: VideoTensorLayout + """Layout of video tensors returned by application sessions.""" + + video_width: int + """Output video width in pixels.""" + + video_height: int + """Output video height in pixels.""" + + default_steps: int | None = None + """Default finite-mode iteration count; ``None`` requires a mode override.""" + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("AppConfig.model_id must be non-empty.") + if float(self.fps) <= 0: + raise ValueError("AppConfig.fps must be > 0.") + if self.video_width <= 0 or self.video_height <= 0: + raise ValueError("AppConfig video dimensions must be > 0.") + if self.default_steps is not None and self.default_steps <= 0: + raise ValueError("AppConfig.default_steps must be > 0 when set.") + + +@dataclass(kw_only=True, slots=True) +class ApplicationArguments: + """Command-line request passed to an application's runtime factory.""" + + mode: str + """Selected runner I/O mode.""" + + parser: argparse.ArgumentParser + """Parser containing runner and selected-mode options.""" + + argv: Sequence[str] + """Arguments remaining after application and mode selection.""" + + _options: argparse.Namespace | None = field(default=None, init=False, repr=False) + + def parse_args(self) -> argparse.Namespace: + """Parse runner, mode, and application options exactly once.""" + if self._options is None: + self._options = self.parser.parse_args(self.argv) + return self._options + + @property + def options(self) -> argparse.Namespace: + """Return options parsed by the application runtime factory.""" + if self._options is None: + raise RuntimeError( + "Application create_runtime() must call arguments.parse_args()." + ) + return self._options + + +class InputHandler(Protocol): + """Supply initial and per-iteration inputs for one runner-owned session.""" + + def open(self) -> None: + """Prepare input resources for a session.""" + ... + + def initial_input(self) -> InferenceInput: + """Return the input used to construct the application session.""" + ... + + def read(self) -> InferenceInput | None: + """Return the next iteration input, or ``None`` to stop the loop.""" + ... + + def close(self) -> None: + """Release input resources.""" + ... + + +class OutputHandler(Protocol): + """Present or persist outputs from one runner-owned session.""" + + def open(self, config: AppConfig) -> None: + """Prepare output resources for a session.""" + ... + + def write(self, result: StepResult) -> None: + """Consume one generated application output.""" + ... + + def close(self) -> Sequence[OutputArtifact]: + """Finalize output resources and return persistent artifacts.""" + ... + + +class Runtime(ABC): + """Application-owned model weights and process-wide inference state. + + The runner creates one runtime for the process, initializes it with the + selected I/O mode, and creates one or more isolated sessions from it. + """ + + @property + @abstractmethod + def config(self) -> AppConfig: + """Return application configuration required by runner-owned modes.""" + + @abstractmethod + def initialize(self, *, device: str, io_handler: "IOHandler") -> None: + """Perform one-time initialization for the selected device and mode.""" + + @abstractmethod + def create_session(self, initial_input: InferenceInput | None = None) -> "Session": + """Create an isolated application session.""" + + @abstractmethod + def destroy(self) -> None: + """Release model weights and process-wide resources.""" + + # These aliases let shared FlashDreams serving code consume the application + # ABI directly while preserving runner lifecycle names. + def start_session(self, inputs: InferenceInput) -> "Session": + """Create a session through the shared inference-runtime API.""" + return self.create_session(inputs) + + def close(self) -> None: + """Destroy the runtime through the shared inference-runtime API.""" + self.destroy() + + def peek_input_fps(self) -> float: + """Return the input clock rate used by realtime presentation.""" + return float(self.config.fps) + + +class Session(ABC): + """Application-owned state and generation logic for one user session.""" + + @property + @abstractmethod + def step_index(self) -> int: + """Return the index of the next generation iteration.""" + + @property + def steady_output_frame_count(self) -> int | None: + """Return the steady output chunk size when the application knows it.""" + return None + + @abstractmethod + def step(self, inputs: InferenceInput) -> StepResult: + """Run one application main-loop iteration.""" + + @abstractmethod + def destroy(self) -> None: + """Release per-session state.""" + + # Shared serving uses the inference-session spelling of this same ABI. + def next_step_request(self) -> StepRequest | None: + """Describe the next iteration, or stop a finite shared session.""" + metadata: dict[str, int] = {} + if self.steady_output_frame_count is not None: + metadata["steady_output_frame_count"] = self.steady_output_frame_count + return StepRequest(step_index=self.step_index, metadata=metadata) + + def reset(self, inputs: InferenceInput | None = None) -> None: + """Reject reset when an application requires a fresh session.""" + del inputs + raise RuntimeError("Create a new application session instead of resetting.") + + def close(self) -> None: + """Destroy the session through the shared inference-session API.""" + self.destroy() + + def session_info(self) -> "SessionInfo": + """Return output information to shared FlashDreams drivers.""" + from flashdreams.runtime.demo import SessionInfo + + return SessionInfo( + steady_output_frame_count=self.steady_output_frame_count, + ) + + +DriveSession = Callable[ + [Runtime, InputHandler, OutputHandler], tuple[OutputArtifact, ...] +] +"""Runner-owned function that drives one application session.""" + + +@runtime_checkable +class IOHandler(Protocol): + """Runner mode that owns input acquisition and output presentation.""" + + @property + def name(self) -> str: + """Return the stable command-line mode name.""" + ... + + def run( + self, + runtime: Runtime, + drive_session: DriveSession, + ) -> tuple[OutputArtifact, ...]: + """Run the mode with an initialized application runtime.""" + ... + + +@runtime_checkable +class Application(Protocol): + """ABI exposed by an installed FlashDreams application module.""" + + def create_runtime(self, arguments: ApplicationArguments) -> Runtime: + """Parse application arguments and return an uninitialized runtime.""" + ... + + +__all__ = [ + "AppConfig", + "Application", + "ApplicationArguments", + "DriveSession", + "IOHandler", + "InputHandler", + "OutputHandler", + "Runtime", + "Session", +] diff --git a/flashdreams_runner/modes.py b/flashdreams_runner/modes.py new file mode 100644 index 000000000..d546b4aa5 --- /dev/null +++ b/flashdreams_runner/modes.py @@ -0,0 +1,172 @@ +# 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. + +"""Runner-owned I/O mode selection and batch mode implementations.""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path + +from flashdreams.runtime import OutputArtifact + +from .contracts import DriveSession, IOHandler, Runtime +from .outputs import FileOutput, FiniteInput, NullOutput +from .webrtc import WebRTCMode + +MODE_MP4 = "mp4" +"""Compatibility name for a replay written to an MP4 file.""" + +MODE_REPLAY = "replay" +"""Finite replay mode that writes an MP4 file.""" + +MODE_WEBRTC = "webrtc" +"""Live WebRTC serving mode.""" + +MODE_NONE = "none" +"""Finite headless mode that discards generated output.""" + +MODE_NAMES = (MODE_MP4, MODE_REPLAY, MODE_WEBRTC, MODE_NONE) +"""I/O modes currently implemented by the application runner.""" + + +@dataclass(frozen=True, slots=True) +class ReplayMode: + """Drive a finite session and persist its output as MP4.""" + + output: Path + """Destination MP4 path.""" + + steps: int | None + """Iteration override; ``None`` uses the application's default.""" + + enabled: bool = True + """Whether this process writes the output artifact.""" + + name: str = MODE_REPLAY + """Stable mode name.""" + + def run( + self, + runtime: Runtime, + drive_session: DriveSession, + ) -> tuple[OutputArtifact, ...]: + """Run one finite application session through MP4 handlers.""" + total_steps = _resolve_steps(self.steps, runtime) + return drive_session( + runtime, + FiniteInput(total_steps=total_steps), + FileOutput(path=self.output, enabled=self.enabled), + ) + + +@dataclass(frozen=True, slots=True) +class NoneMode: + """Drive a finite session while discarding all generated output.""" + + steps: int | None + """Iteration override; ``None`` uses the application's default.""" + + name: str = MODE_NONE + """Stable mode name.""" + + def run( + self, + runtime: Runtime, + drive_session: DriveSession, + ) -> tuple[OutputArtifact, ...]: + """Run one finite application session through headless handlers.""" + total_steps = _resolve_steps(self.steps, runtime) + return drive_session( + runtime, + FiniteInput(total_steps=total_steps), + NullOutput(), + ) + + +def add_mode_arguments(parser: argparse.ArgumentParser, mode: str) -> None: + """Add runner-owned arguments for the selected I/O mode.""" + parser.add_argument("--device", default="cuda", help="Runtime device") + if mode in (MODE_MP4, MODE_REPLAY): + parser.add_argument("--output", type=Path, required=True, help="MP4 path") + parser.add_argument( + "--steps", + type=int, + help="Generation iterations (defaults to the application preset)", + ) + return + if mode == MODE_NONE: + parser.add_argument( + "--steps", + type=int, + help="Generation iterations (defaults to the application preset)", + ) + return + if mode == MODE_WEBRTC: + parser.add_argument("--host", default="0.0.0.0", help="WebRTC bind address") + parser.add_argument("--port", type=int, default=8080, help="WebRTC bind port") + return + raise ValueError(f"Unsupported application mode: {mode!r}.") + + +def create_io_handler( + mode: str, + options: argparse.Namespace, + *, + device: str, + world_rank: int, +) -> IOHandler: + """Create the selected runner-owned I/O handler.""" + if mode in (MODE_MP4, MODE_REPLAY): + return ReplayMode( + output=options.output, + steps=options.steps, + enabled=world_rank == 0, + name=mode, + ) + if mode == MODE_NONE: + return NoneMode(steps=options.steps) + if mode == MODE_WEBRTC: + return WebRTCMode( + host=options.host, + port=options.port, + device=device, + world_rank=world_rank, + ) + raise ValueError(f"Unsupported application mode: {mode!r}.") + + +def _resolve_steps(value: int | None, runtime: Runtime) -> int: + total_steps = runtime.config.default_steps if value is None else value + if total_steps is None: + raise ValueError( + "Finite modes require --steps or an application default step count." + ) + if total_steps <= 0: + raise ValueError("steps must be > 0.") + return total_steps + + +__all__ = [ + "MODE_NAMES", + "MODE_NONE", + "MODE_REPLAY", + "MODE_WEBRTC", + "NoneMode", + "ReplayMode", + "add_mode_arguments", + "create_io_handler", +] diff --git a/flashdreams_runner/outputs.py b/flashdreams_runner/outputs.py new file mode 100644 index 000000000..73e6ba490 --- /dev/null +++ b/flashdreams_runner/outputs.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Batch input and output handlers for runner modes.""" + +from __future__ import annotations + +from pathlib import Path + +from flashdreams.runtime import ( + InferenceInput, + NullOutputTarget, + OutputArtifact, + StepResult, +) +from flashdreams.runtime.video_output import Mp4VideoOutputTarget + +from .contracts import AppConfig + + +class FiniteInput: + """Supply empty per-step inputs for a fixed number of iterations.""" + + def __init__(self, *, total_steps: int) -> None: + if total_steps <= 0: + raise ValueError("FiniteInput.total_steps must be > 0.") + self._total_steps = total_steps + self._step_index = 0 + self._opened = False + + def open(self) -> None: + """Reset the finite input sequence.""" + self._step_index = 0 + self._opened = True + + def initial_input(self) -> InferenceInput: + """Return an empty initial input for application-owned defaults.""" + if not self._opened: + raise RuntimeError("Cannot read from a closed input handler.") + return InferenceInput() + + def read(self) -> InferenceInput | None: + """Return one empty iteration input until the configured limit.""" + if not self._opened: + raise RuntimeError("Cannot read from a closed input handler.") + if self._step_index >= self._total_steps: + return None + self._step_index += 1 + return InferenceInput() + + def close(self) -> None: + """Close the finite input sequence.""" + self._opened = False + + +class FileOutput: + """Collect generated chunks and write one MP4 file.""" + + def __init__(self, *, path: Path, enabled: bool = True) -> None: + self._path = path + self._enabled = enabled + self._target: Mp4VideoOutputTarget | None = None + + def open(self, config: AppConfig) -> None: + """Open a video target using application output information.""" + if self._target is not None: + raise RuntimeError("FileOutput is already open.") + self._target = Mp4VideoOutputTarget( + output_path=self._path, + fps=config.fps, + output_layout=config.output_layout, + enabled=self._enabled, + ) + self._target.open() + + def write(self, result: StepResult) -> None: + """Append one generated output chunk.""" + if self._target is None: + raise RuntimeError("Cannot write to a closed FileOutput.") + self._target.write(result) + + def close(self) -> tuple[OutputArtifact, ...]: + """Finalize the MP4 and return its artifact metadata.""" + if self._target is None: + return () + target = self._target + self._target = None + return tuple(target.close()) + + +class NullOutput: + """Discard generated outputs for the ``none`` mode.""" + + def __init__(self) -> None: + self._target = NullOutputTarget() + + def open(self, config: AppConfig) -> None: + """Open the headless output target.""" + del config + self._target.open() + + def write(self, result: StepResult) -> None: + """Discard one generated output while recording its count.""" + self._target.write(result) + + def close(self) -> tuple[OutputArtifact, ...]: + """Close the headless output target without creating artifacts.""" + return tuple(self._target.close()) + + +__all__ = ["FileOutput", "FiniteInput", "NullOutput"] diff --git a/flashdreams_runner/pyproject.toml b/flashdreams_runner/pyproject.toml new file mode 100644 index 000000000..16690f24a --- /dev/null +++ b/flashdreams_runner/pyproject.toml @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-runner" +version = "0.1.0" +description = "Generic shell for FlashDreams applications and I/O modes" +requires-python = ">=3.10" +dependencies = ["flashdreams[serving]"] + +[project.scripts] +flashdreams-runner = "flashdreams_runner.cli:main" + +[tool.uv.sources] +flashdreams = { workspace = true } + +[tool.setuptools] +packages = ["flashdreams_runner"] +package-dir = { flashdreams_runner = "." } diff --git a/flashdreams_runner/tests/test_cli.py b/flashdreams_runner/tests/test_cli.py new file mode 100644 index 000000000..a8980f3ae --- /dev/null +++ b/flashdreams_runner/tests/test_cli.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the application runner lifecycle and mode boundary.""" + +from __future__ import annotations + +import argparse +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +import flashdreams_runner +from flashdreams.runtime import InferenceInput, OutputArtifact, StepResult +from flashdreams_runner import ( + AppConfig, + Application, + ApplicationArguments, + DriveSession, + IOHandler, + Runtime, + Session, + cli, +) + +pytestmark = pytest.mark.ci_cpu + + +def _config() -> AppConfig: + return AppConfig( + model_id="fake-app", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + default_steps=1, + ) + + +def _result(index: int = 0) -> StepResult: + return StepResult.from_video_chunk( + step_index=index, + video_chunk=torch.zeros((1, 3, 2, 2)), + layout="tchw", + ) + + +def test_public_package_surface_is_the_application_abi() -> None: + assert flashdreams_runner.__all__ == [ + "AppConfig", + "Application", + "ApplicationArguments", + "DriveSession", + "IOHandler", + "InputHandler", + "OutputHandler", + "Runtime", + "Session", + ] + + +def test_runner_owns_lifecycle_io_and_main_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + class FakeSession(Session): + def __init__(self) -> None: + self._step_index = 0 + + @property + def step_index(self) -> int: + return self._step_index + + def step(self, inputs: InferenceInput) -> StepResult: + assert not inputs.global_conditioning + calls.append("session.step") + result = _result(self._step_index) + self._step_index += 1 + return result + + def destroy(self) -> None: + calls.append("session.destroy") + + class FakeRuntime(Runtime): + @property + def config(self) -> AppConfig: + return _config() + + def initialize(self, *, device: str, io_handler: IOHandler) -> None: + assert device == "cpu" + assert io_handler.name == "fake" + calls.append("runtime.initialize") + + def create_session( + self, initial_input: InferenceInput | None = None + ) -> Session: + assert isinstance(initial_input, InferenceInput) + calls.append("runtime.create_session") + return FakeSession() + + def destroy(self) -> None: + calls.append("runtime.destroy") + + runtime = FakeRuntime() + application = ModuleType("fake_app") + + def create_runtime(arguments: ApplicationArguments) -> Runtime: + calls.append("application.create_runtime") + arguments.parser.add_argument("--model-option", required=True) + options = arguments.parse_args() + assert options.model_option == "enabled" + return runtime + + setattr(application, "create_runtime", create_runtime) + monkeypatch.setattr(cli, "load_application", lambda _: application) + + class Input: + def open(self) -> None: + calls.append("input.open") + + def initial_input(self) -> InferenceInput: + calls.append("input.initial_input") + return InferenceInput() + + def read(self) -> InferenceInput | None: + calls.append("input.read") + if calls.count("input.read") == 1: + return InferenceInput() + return None + + def close(self) -> None: + calls.append("input.close") + + class Output: + def open(self, config: AppConfig) -> None: + assert config.model_id == "fake-app" + calls.append("output.open") + + def write(self, result: StepResult) -> None: + assert result.step_index == 0 + calls.append("output.write") + + def close(self) -> tuple[OutputArtifact, ...]: + calls.append("output.close") + return () + + class Mode: + name = "fake" + + def run( + self, runtime: Runtime, drive_session: DriveSession + ) -> tuple[OutputArtifact, ...]: + calls.append("mode.run") + return drive_session(runtime, Input(), Output()) + + mode = Mode() + monkeypatch.setattr(cli, "create_io_handler", lambda *args, **kwargs: mode) + + assert ( + cli.run( + [ + "fake-app", + "mp4", + "--device", + "cpu", + "--output", + "result.mp4", + "--model-option", + "enabled", + ] + ) + == () + ) + assert calls == [ + "application.create_runtime", + "runtime.initialize", + "mode.run", + "input.open", + "output.open", + "input.initial_input", + "runtime.create_session", + "input.read", + "session.step", + "output.write", + "input.read", + "output.close", + "session.destroy", + "input.close", + "runtime.destroy", + ] + + +def test_application_protocol_requires_only_runtime_factory() -> None: + application = ModuleType("application") + setattr(application, "create_runtime", lambda arguments: None) + assert isinstance(application, Application) + + delattr(application, "create_runtime") + assert not isinstance(application, Application) + + +def test_load_application_rejects_module_outside_abi( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = ModuleType("invalid_application") + distribution = SimpleNamespace(metadata={"Name": "invalid-app"}) + monkeypatch.setattr(cli.metadata, "distribution", lambda name: distribution) + monkeypatch.setattr( + cli.metadata, + "packages_distributions", + lambda: {"invalid_application": ["invalid-app"]}, + ) + monkeypatch.setattr(cli.importlib, "import_module", lambda name: module) + + with pytest.raises(TypeError, match="none satisfy Application"): + cli.load_application("invalid-app") + + +def test_runner_exposes_mode_names_and_preserves_application_arguments() -> None: + route = cli._parse_application_and_mode(["fake-app", "webrtc", "--prompt", "x"]) + assert route.application == "fake-app" + assert route.mode == "webrtc" + assert route.remaining_argv == ("--prompt", "x") + + for mode in ("mp4", "replay", "webrtc", "none"): + assert cli._parse_application_and_mode(["fake-app", mode]).mode == mode + + with pytest.raises(SystemExit): + cli._parse_application_and_mode(["fake-app", "unsupported"]) + + +def test_mode_parsers_keep_transport_options_separate() -> None: + mp4_destinations = { + action.dest for action in cli.build_parser("fake-app", "mp4")._actions + } + assert {"device", "output", "steps"} <= mp4_destinations + assert {"host", "port"}.isdisjoint(mp4_destinations) + + webrtc_destinations = { + action.dest for action in cli.build_parser("fake-app", "webrtc")._actions + } + assert {"device", "host", "port"} <= webrtc_destinations + assert {"output", "steps"}.isdisjoint(webrtc_destinations) + + +def test_application_arguments_must_be_parsed_by_factory() -> None: + arguments = ApplicationArguments( + mode="none", + parser=argparse.ArgumentParser(), + argv=(), + ) + with pytest.raises(RuntimeError, match="must call arguments.parse_args"): + _ = arguments.options + + assert isinstance(arguments.parse_args(), argparse.Namespace) + assert arguments.options is arguments.parse_args() + + +def test_runtime_and_session_bridge_shared_inference_api() -> None: + calls: list[str] = [] + + class FakeSession(Session): + @property + def step_index(self) -> int: + return 3 + + def step(self, inputs: InferenceInput) -> StepResult: + calls.append("step") + return _result(3) + + def destroy(self) -> None: + calls.append("session.destroy") + + session = FakeSession() + + class FakeRuntime(Runtime): + @property + def config(self) -> AppConfig: + return _config() + + def initialize(self, *, device: str, io_handler: IOHandler) -> None: + del device, io_handler + + def create_session( + self, initial_input: InferenceInput | None = None + ) -> Session: + calls.append("create_session") + return session + + def destroy(self) -> None: + calls.append("runtime.destroy") + + runtime = FakeRuntime() + assert runtime.start_session(InferenceInput()) is session + request = session.next_step_request() + assert request is not None + assert request.step_index == 3 + assert session.step(InferenceInput()).step_index == 3 + session.close() + runtime.close() + assert calls == [ + "create_session", + "step", + "session.destroy", + "runtime.destroy", + ] diff --git a/flashdreams_runner/tests/test_webrtc.py b/flashdreams_runner/tests/test_webrtc.py new file mode 100644 index 000000000..289dbd095 --- /dev/null +++ b/flashdreams_runner/tests/test_webrtc.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the WebRTC application mode.""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest + +from flashdreams.runtime import InferenceInput, StepResult +from flashdreams.runtime.demo import WebRTCAppResources +from flashdreams_runner import AppConfig, IOHandler, Runtime, Session, webrtc + +pytestmark = pytest.mark.ci_cpu + + +class _Session(Session): + @property + def step_index(self) -> int: + return 0 + + def step(self, inputs: InferenceInput) -> StepResult: + del inputs + raise AssertionError("Server construction must not run a step.") + + def destroy(self) -> None: + pass + + +class _Runtime(Runtime): + @property + def config(self) -> AppConfig: + return AppConfig( + model_id="fake-app", + fps=16, + output_layout="tchw", + video_width=96, + video_height=64, + ) + + def initialize(self, *, device: str, io_handler: IOHandler) -> None: + del device, io_handler + + def create_session(self, initial_input: InferenceInput | None = None) -> Session: + del initial_input + return _Session() + + def destroy(self) -> None: + pass + + +def test_webrtc_mode_constructs_shared_presentation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def fake_serve(**kwargs: object) -> str: + captured.update(kwargs) + return "served" + + monkeypatch.setattr(webrtc, "serve_webrtc_demo", fake_serve) + runtime = _Runtime() + result = webrtc.serve_webrtc( + runtime=runtime, + host="127.0.0.1", + port=8080, + device="cpu", + world_rank=0, + ) + + assert result == "served" + assert captured["model_id"] == "fake-app" + assert captured["world_rank"] == 0 + output = captured["output"] + assert isinstance(output, webrtc.WebRTCOutputSpec) + assert output.video_width == 96 + assert output.warmup_chunks == 0 + session_manager = captured["session_manager"] + assert isinstance(session_manager, webrtc.BaseWebRTCSessionManager) + assert session_manager._shared_adapter is None + assert session_manager._shared_host is not None + assert session_manager._shared_host.runtime is runtime + assert session_manager.is_runtime_ready() + + +def test_webrtc_mode_uses_application_customization( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + custom_manager = object() + + class Customization: + def prepare_initial_input(self) -> InferenceInput: + return InferenceInput(global_conditioning={"prompt": "custom"}) + + def create_session_manager(self, **kwargs: object) -> Any: + captured["manager_kwargs"] = kwargs + return custom_manager + + def create_app_resources(self, **kwargs: object) -> WebRTCAppResources: + captured["resources_kwargs"] = kwargs + return WebRTCAppResources(preload_name="custom-ui") + + def fake_serve(**kwargs: object) -> str: + captured.update(kwargs) + return "served" + + monkeypatch.setattr(webrtc, "serve_webrtc_demo", fake_serve) + result = webrtc.serve_webrtc( + runtime=_Runtime(), + host="127.0.0.1", + port=8080, + device="cpu", + world_rank=0, + customization=Customization(), + ) + + assert result == "served" + assert captured["session_manager"] is custom_manager + resources = captured["app_resources"] + assert isinstance(resources, WebRTCAppResources) + assert resources.preload_name == "custom-ui" + manager_kwargs = cast(dict[str, object], captured["manager_kwargs"]) + scenario = cast(Any, manager_kwargs["scenario"]) + assert scenario.initial_inputs.global_conditioning["prompt"] == "custom" + assert ( + cast(dict[str, object], captured["resources_kwargs"])["session_manager"] + is custom_manager + ) diff --git a/flashdreams_runner/webrtc.py b/flashdreams_runner/webrtc.py new file mode 100644 index 000000000..661be7178 --- /dev/null +++ b/flashdreams_runner/webrtc.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""WebRTC I/O mode for application runtimes.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Protocol, cast + +from flashdreams.runtime import ( + InferenceConfig, + InferenceInput, + InferenceRuntime, + OutputArtifact, +) +from flashdreams.runtime.demo import ( + DemoSpec, + ModelInputProvider, + PreparedScenario, + PreparedStep, + ProviderCapabilities, + RuntimeHost, + UserInputWindow, + WebRTCAppResources, + WebRTCOutputSpec, +) +from flashdreams.runtime.types import StepRequirements +from flashdreams.serving.webrtc.demo import serve_webrtc_demo +from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager +from flashdreams.serving.webrtc.runtime import WebRTCRuntimeConfig + +from .contracts import DriveSession, Runtime + + +@dataclass(slots=True) +class WebRTCMode: + """Serve application sessions over WebRTC.""" + + host: str + """Server bind address.""" + + port: int + """Server bind port.""" + + device: str + """Device used by the application runtime.""" + + world_rank: int + """Distributed rank responsible for presentation.""" + + name: str = "webrtc" + """Stable mode name.""" + + _customization: WebRTCCustomization | None = field( + default=None, + init=False, + repr=False, + ) + + def customize(self, customization: WebRTCCustomization) -> None: + """Install application-owned WebRTC behavior before serving starts.""" + if self._customization is not None: + raise RuntimeError("WebRTC mode customization is already installed.") + self._customization = customization + + def run( + self, + runtime: Runtime, + drive_session: DriveSession, + ) -> tuple[OutputArtifact, ...]: + """Serve live sessions until the WebRTC server exits.""" + del drive_session + serve_webrtc( + runtime=runtime, + host=self.host, + port=self.port, + device=self.device, + world_rank=self.world_rank, + customization=self._customization, + ) + return () + + +class _InputProvider: + """Provide prepared session input and empty per-step transport input.""" + + capabilities = ProviderCapabilities( + supports_realtime_clock=True, + deterministic_given_inputs=True, + ) + + def __init__(self, initial_input: InferenceInput) -> None: + self._initial_input = initial_input + + def prepare_initial_input(self) -> InferenceInput: + """Return the application-provided input for a new session.""" + return self._initial_input + + def prepare_step( + self, *, request: StepRequirements, user_window: UserInputWindow + ) -> PreparedStep: + """Return one empty step input for a non-interactive T2V session.""" + del request, user_window + return PreparedStep(inference_input=InferenceInput()) + + def reset(self, inputs: InferenceInput | None = None) -> None: + """Replace the initial input when the shared driver requests a reset.""" + if inputs is not None: + self._initial_input = inputs + + def close(self) -> None: + """Release provider resources.""" + + +ModelInputProviderFactory = Callable[ + [DemoSpec, PreparedScenario], + ModelInputProvider, +] +"""Factory used by shared WebRTC drivers to prepare application input.""" + + +class WebRTCCustomization(Protocol): + """Application-owned WebRTC UI and session-manager extension point. + + Applications only implement this interface when the generic WebRTC viewer + is insufficient. The runner continues to own the server and transport. + """ + + def prepare_initial_input(self) -> InferenceInput: + """Return the initial input used for the first browser generation.""" + ... + + def create_session_manager( + self, + *, + runtime: Runtime, + output: WebRTCOutputSpec, + spec: DemoSpec, + scenario: PreparedScenario, + input_provider_factory: ModelInputProviderFactory, + ) -> BaseWebRTCSessionManager[Any, Any]: + """Create the transport manager used by the customized application.""" + ... + + def create_app_resources( + self, + *, + session_manager: BaseWebRTCSessionManager[Any, Any], + ) -> WebRTCAppResources: + """Return packaged browser assets and optional HTTP routes.""" + ... + + +def serve_webrtc( + *, + runtime: Runtime, + host: str, + port: int, + device: str, + world_rank: int, + customization: WebRTCCustomization | None = None, +) -> object: + """Serve an initialized application runtime through shared WebRTC. + + Args: + runtime: Initialized application runtime. + host: Server bind address. + port: Server bind port. + device: Device used by the runtime. + world_rank: Distributed rank responsible for presentation. + customization: Optional application-owned UI and manager behavior. + + Returns: + Serving backend result. + """ + config = runtime.config + output = WebRTCOutputSpec( + host=host, + port=port, + fps=int(config.fps), + video_width=config.video_width, + video_height=config.video_height, + preload_name=config.model_id, + ) + spec = DemoSpec( + model_id=config.model_id, + input_mode="webrtc", + output=output, + config=InferenceConfig(model_id=config.model_id, device=device), + ) + initial_input = ( + InferenceInput() + if customization is None + else customization.prepare_initial_input() + ) + scenario = PreparedScenario(initial_inputs=initial_input) + + def create_model_input_provider( + spec: DemoSpec, + scenario: PreparedScenario, + ) -> _InputProvider: + del spec + return _InputProvider(scenario.initial_inputs) + + inference_runtime = cast(InferenceRuntime, runtime) + if customization is None: + manager: BaseWebRTCSessionManager[Any, Any] = BaseWebRTCSessionManager( + runtime=inference_runtime, + runtime_config=cast(WebRTCRuntimeConfig, cast(object, output)), + fps=int(config.fps), + identity=config.model_id, + shared_host=RuntimeHost(inference_runtime), + shared_spec=spec, + shared_scenario=scenario, + shared_model_input_provider_factory=create_model_input_provider, + client_liveness_timeout_s=output.client_liveness_timeout_s, + runtime_ready=True, + ) + app_resources = WebRTCAppResources(preload_name=config.model_id) + else: + manager = customization.create_session_manager( + runtime=runtime, + output=output, + spec=spec, + scenario=scenario, + input_provider_factory=create_model_input_provider, + ) + app_resources = customization.create_app_resources( + session_manager=manager, + ) + return serve_webrtc_demo( + output=output, + model_id=config.model_id, + session_manager=manager, + app_resources=app_resources, + world_rank=world_rank, + ) + + +__all__ = [ + "ModelInputProviderFactory", + "WebRTCCustomization", + "WebRTCMode", + "serve_webrtc", +] diff --git a/integrations/causal_forcing/causal_forcing/config.py b/integrations/causal_forcing/causal_forcing/config.py index 98242d8d9..b84781f79 100644 --- a/integrations/causal_forcing/causal_forcing/config.py +++ b/integrations/causal_forcing/causal_forcing/config.py @@ -17,14 +17,13 @@ from __future__ import annotations -from typing import Any, cast - -from torch import Tensor +from typing import cast from causal_forcing.runner import ( CausalForcingI2VRunnerConfig, CausalForcingT2VRunnerConfig, ) +from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig @@ -42,30 +41,8 @@ CHECKPOINT_PATH_FRAMEWISE = "https://huggingface.co/zhuhz22/Causal-Forcing/blob/main/framewise/causal_forcing.pt" -def state_dict_transform(state_dict: dict[str, Any]) -> dict[str, Tensor]: - """Strip Causal-Forcing wrapper prefixes from the checkpoint state-dict. - - Drops the ``generator_ema`` / ``generator`` container, the ``model.`` - / ``net.`` outer prefix, and the ``_fsdp_wrapped_module.`` inner - prefix (framewise variant) so keys match a bare ``WanDiTNetwork``. - """ - if "generator_ema" in state_dict: - state_dict = state_dict["generator_ema"] - elif "generator" in state_dict: - state_dict = state_dict["generator"] - - out: dict[str, Tensor] = {} - for k, v in state_dict.items(): - if k.startswith("model."): - new_k = k[len("model.") :] - elif k.startswith("net."): - new_k = k[len("net.") :] - else: - new_k = k - if new_k.startswith("_fsdp_wrapped_module."): - new_k = new_k[len("_fsdp_wrapped_module.") :] - out[new_k] = v - return out +state_dict_transform = unwrap_generator_state_dict +"""State-dict transform for Causal-Forcing generator checkpoint envelopes.""" # Causal-Forcing chunkwise Wan 2.1 1.3B T2V pipeline. diff --git a/integrations/self_forcing/self_forcing/config.py b/integrations/self_forcing/self_forcing/config.py index becf1ecd5..3d02ed332 100644 --- a/integrations/self_forcing/self_forcing/config.py +++ b/integrations/self_forcing/self_forcing/config.py @@ -17,10 +17,9 @@ from __future__ import annotations -from typing import Any, cast - -from torch import Tensor +from typing import cast +from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict from flashdreams.infra.config import derive_config from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig @@ -37,25 +36,8 @@ CHECKPOINT_PATH = "https://huggingface.co/gdhe17/Self-Forcing/blob/main/checkpoints/self_forcing_dmd.pt" -def state_dict_transform(state_dict: dict[str, Any]) -> dict[str, Tensor]: - """Strip Self-Forcing wrapper prefixes from the checkpoint state-dict.""" - if "generator_ema" in state_dict: - state_dict = state_dict["generator_ema"] - elif "generator" in state_dict: - state_dict = state_dict["generator"] - - out: dict[str, Tensor] = {} - for k, v in state_dict.items(): - if k.startswith("model."): - new_k = k[len("model.") :] - elif k.startswith("net."): - new_k = k[len("net.") :] - else: - new_k = k - if new_k.startswith("_fsdp_wrapped_module."): - new_k = new_k[len("_fsdp_wrapped_module.") :] - out[new_k] = v - return out +state_dict_transform = unwrap_generator_state_dict +"""State-dict transform for Self-Forcing generator checkpoint envelopes.""" # Official Self-Forcing Wan 2.1 1.3B T2V pipeline config. diff --git a/pyproject.toml b/pyproject.toml index 02f62aa87..916eea902 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ members = [ # extractions and don't ship a ``pyproject.toml`` yet. "integrations/*", "apps/*", + "flashdreams_runner", # Nested sub-packages that the ``integrations/*`` glob does not reach. "integrations/omnidreams/ludus-renderer", ] @@ -45,6 +46,7 @@ no-build-isolation-package = ["transformer-engine-torch"] extraPaths = [ "flashdreams", "apps", + "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", "integrations/causal_forcing", @@ -70,6 +72,7 @@ python-version = "3.10" extra-paths = [ "flashdreams", "apps", + "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", "integrations/causal_forcing", diff --git a/uv.lock b/uv.lock index 7f4aba302..af4874da6 100644 --- a/uv.lock +++ b/uv.lock @@ -27,12 +27,14 @@ members = [ "flashdreams-hy-worldplay", "flashdreams-lingbot", "flashdreams-omnidreams", + "flashdreams-runner", "flashdreams-sana-wm", "flashdreams-self-forcing", "flashdreams-t2v-demo", "flashdreams-wan21", "flashdreams-wan22", "ludus-renderer", + "t2v-app", ] overrides = [ { name = "numpy", specifier = ">=1.24,<2.5" }, @@ -1356,6 +1358,17 @@ requires-dist = [ ] provides-extras = ["interactive-drive", "rtx-postprocess", "dev"] +[[package]] +name = "flashdreams-runner" +version = "0.1.0" +source = { editable = "flashdreams_runner" } +dependencies = [ + { name = "flashdreams", extra = ["serving"] }, +] + +[package.metadata] +requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] + [[package]] name = "flashdreams-sana-wm" version = "0.1.0" @@ -4502,6 +4515,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "t2v-app" +version = "0.1.0" +source = { editable = "apps/t2v_app" } +dependencies = [ + { name = "flashdreams" }, + { name = "flashdreams-runner" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-runner", editable = "flashdreams_runner" }, +] + [[package]] name = "tokenizers" version = "0.22.2"