From 387de18f3330d588a7964f35a34393dff5299281 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 02:19:09 +0000 Subject: [PATCH 01/10] Add modular FlashDreams application runner Signed-off-by: Gangzheng Tong --- apps/flashdreams_app/README.md | 26 + .../flashdreams_app/__init__.py | 24 + apps/flashdreams_app/flashdreams_app/cli.py | 194 ++++++++ .../flashdreams_app/contracts.py | 132 +++++ .../flashdreams_app/outputs.py | 44 ++ .../flashdreams_app/runtime.py | 197 ++++++++ .../flashdreams_app/flashdreams_app/webrtc.py | 138 ++++++ apps/flashdreams_app/pyproject.toml | 22 + apps/flashdreams_app/tests/test_cli.py | 190 ++++++++ apps/flashdreams_app/tests/test_webrtc.py | 87 ++++ apps/t2v_app/README.md | 50 ++ apps/t2v_app/pyproject.toml | 24 + apps/t2v_app/t2v_app/__init__.py | 10 + apps/t2v_app/t2v_app/pipeline_presets.yaml | 459 ++++++++++++++++++ apps/t2v_app/t2v_app/presets.py | 159 ++++++ apps/t2v_app/t2v_app/provider.py | 203 ++++++++ apps/t2v_app/tests/test_pipeline_provider.py | 88 ++++ apps/t2v_app/tests/test_provider.py | 101 ++++ .../flashdreams/core/checkpoint/remap.py | 39 +- .../flashdreams/core/pipeline_presets.py | 357 ++++++++++++++ .../flashdreams/runtime/demo/drivers.py | 14 +- .../flashdreams/serving/webrtc/manager.py | 37 +- flashdreams/tests/test_checkpoint_loading.py | 46 ++ .../tests/test_demo_runtime_run_modes.py | 61 +++ flashdreams/tests/test_pipeline_presets.py | 117 +++++ flashdreams/tests/test_webrtc_manager.py | 67 ++- .../causal_forcing/causal_forcing/config.py | 31 +- .../self_forcing/self_forcing/config.py | 26 +- pyproject.toml | 4 + uv.lock | 28 ++ 30 files changed, 2915 insertions(+), 60 deletions(-) create mode 100644 apps/flashdreams_app/README.md create mode 100644 apps/flashdreams_app/flashdreams_app/__init__.py create mode 100644 apps/flashdreams_app/flashdreams_app/cli.py create mode 100644 apps/flashdreams_app/flashdreams_app/contracts.py create mode 100644 apps/flashdreams_app/flashdreams_app/outputs.py create mode 100644 apps/flashdreams_app/flashdreams_app/runtime.py create mode 100644 apps/flashdreams_app/flashdreams_app/webrtc.py create mode 100644 apps/flashdreams_app/pyproject.toml create mode 100644 apps/flashdreams_app/tests/test_cli.py create mode 100644 apps/flashdreams_app/tests/test_webrtc.py create mode 100644 apps/t2v_app/README.md create mode 100644 apps/t2v_app/pyproject.toml create mode 100644 apps/t2v_app/t2v_app/__init__.py create mode 100644 apps/t2v_app/t2v_app/pipeline_presets.yaml create mode 100644 apps/t2v_app/t2v_app/presets.py create mode 100644 apps/t2v_app/t2v_app/provider.py create mode 100644 apps/t2v_app/tests/test_pipeline_provider.py create mode 100644 apps/t2v_app/tests/test_provider.py create mode 100644 flashdreams/flashdreams/core/pipeline_presets.py create mode 100644 flashdreams/tests/test_pipeline_presets.py diff --git a/apps/flashdreams_app/README.md b/apps/flashdreams_app/README.md new file mode 100644 index 000000000..b1ca3d872 --- /dev/null +++ b/apps/flashdreams_app/README.md @@ -0,0 +1,26 @@ +# FlashDreams App Host + +`flashdreams-app` is the model-neutral application host. It finds a provider +distribution in the active Python environment, asks it for a declarative +pipeline application spec, constructs the runtime, creates a session, runs the +session loop, and writes presentation artifacts itself. + +```bash +uv run flashdreams-app t2v-app mp4 --output o.mp4 --prompt "A waterfall" +uv run flashdreams-app t2v-app webrtc --prompt "A waterfall" +``` + +## Provider contract + +A compatible package must expose an importable module with: + +- `create_app(config: flashdreams_app.AppConfig)`, returning a + `PipelineAppSpec` with a `StreamInferencePipelineConfig`, initial + conditioning, presentation metadata, step count, and a `PipelineContract` + cache initializer. +- Optionally, `add_arguments(parser)` adds provider-specific flags. + +The host owns process/distributed initialization, pipeline setup, execution +options, runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, +and WebRTC serving. Providers only select/configure a pipeline and describe how +global conditioning initializes its cache. diff --git a/apps/flashdreams_app/flashdreams_app/__init__.py b/apps/flashdreams_app/flashdreams_app/__init__.py new file mode 100644 index 000000000..9a17f6ab0 --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/__init__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host runtime and public provider contract for FlashDreams applications.""" + +from .contracts import ( + AppConfig, + AppRuntime, + PipelineAppSpec, + PipelineContract, + RuntimeMetadata, + require_pipeline_config, +) +from .runtime import PipelineAppRuntime + +__all__ = [ + "AppConfig", + "AppRuntime", + "PipelineAppRuntime", + "PipelineAppSpec", + "PipelineContract", + "RuntimeMetadata", + "require_pipeline_config", +] diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py new file mode 100644 index 000000000..8058e5414 --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command line host for independently installed FlashDreams app providers.""" + +from __future__ import annotations + +import argparse +import importlib +from contextlib import ExitStack +from importlib import metadata +from pathlib import Path +from types import ModuleType +from typing import Sequence + +import torch + +from flashdreams.runtime.demo.bootstrap import ( + configure_logging, + initialize_cuda_distributed, +) +from flashdreams.runtime.output import OutputArtifact + +from .contracts import AppConfig, PipelineAppSpec +from .outputs import FileOutput +from .runtime import PipelineAppRuntime +from .webrtc import WebRTCOptions, serve_webrtc + + +def build_parser() -> argparse.ArgumentParser: + """Build the host parser without importing a provider package.""" + parser = argparse.ArgumentParser(prog="flashdreams-app") + parser.add_argument( + "provider", help="Installed provider distribution, e.g. t2v-app" + ) + parser.add_argument("mode", choices=("mp4", "webrtc")) + parser.add_argument("--output", type=Path, help="MP4 path (required for mp4)") + parser.add_argument("--device", default="cuda", help="Runtime device") + parser.add_argument( + "--compile", action=argparse.BooleanOptionalAction, default=None + ) + parser.add_argument( + "--cuda-graph", action=argparse.BooleanOptionalAction, default=None + ) + 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") + parser.add_argument("--warmup-chunks", type=int, default=0) + parser.add_argument("--warmup-timeout-s", type=float, default=600.0) + parser.add_argument("--client-liveness-timeout-s", type=float, default=30.0) + parser.add_argument( + "--encoder-backend", choices=("auto", "default", "nvenc"), default="auto" + ) + parser.add_argument("--encoder-bitrate-bps", type=int, default=6_000_000) + parser.add_argument("--encoder-gop", type=int) + return parser + + +def load_provider(distribution_name: str) -> ModuleType: + """Load an installed provider after verifying its distribution is present.""" + try: + distribution = metadata.distribution(distribution_name) + except metadata.PackageNotFoundError as exc: + raise ValueError( + f"Provider 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("-", "_")) + for candidate in dict.fromkeys(candidates): + try: + return importlib.import_module(candidate) + except ModuleNotFoundError as exc: + if exc.name != candidate: + raise + raise ValueError( + f"Provider {distribution_name!r} does not expose an importable Python module." + ) + + +def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: + """Run one provider session and return artifacts produced by the host.""" + probe = argparse.ArgumentParser(add_help=False) + probe.add_argument("provider", nargs="?") + probe.add_argument("mode", nargs="?") + provider_args, _ = probe.parse_known_args(argv) + parser = build_parser() + if provider_args.provider is None: + parser.parse_args(argv) + return () + provider = load_provider(provider_args.provider) + add_arguments = getattr(provider, "add_arguments", None) + if callable(add_arguments): + add_arguments(parser) + args = parser.parse_args(argv) + if args.mode == "mp4" and args.output is None: + parser.error("--output is required for mp4 mode") + + environment = _initialize_environment(args.device) + options = vars(args).copy() + options["device"] = environment.device + options["world_rank"] = environment.world_rank + options["world_size"] = environment.world_size + + factory = getattr(provider, "create_app", None) + if not callable(factory): + raise TypeError(f"Provider {args.provider!r} must define create_app(config).") + spec = factory(AppConfig(options=options)) + if not isinstance(spec, PipelineAppSpec): + raise TypeError( + f"Provider {args.provider!r} create_app() returned " + f"{type(spec).__name__}, expected PipelineAppSpec." + ) + runtime = PipelineAppRuntime( + spec=spec, + device=environment.device, + compile=args.compile, + cuda_graph=args.cuda_graph, + ) + if args.mode == "webrtc": + try: + serve_webrtc( + runtime=runtime, + options=WebRTCOptions( + host=args.host, + port=args.port, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + device=environment.device, + encoder_backend=args.encoder_backend, + encoder_bitrate_bps=args.encoder_bitrate_bps, + encoder_gop=args.encoder_gop or int(runtime.metadata.fps), + ), + world_rank=environment.world_rank, + ) + finally: + runtime.close() + return () + + with ExitStack() as resources: + resources.callback(runtime.close) + output = FileOutput( + path=args.output, + fps=runtime.metadata.fps, + output_layout=runtime.metadata.output_layout, + enabled=environment.world_rank == 0, + ) + output_closed = False + + def close_output() -> None: + if not output_closed: + output.close() + + resources.callback(close_output) + output.open() + session = runtime.start_session(runtime.initial_input) + resources.callback(session.close) + while (request := session.next_step_request()) is not None: + output.write(session.step(runtime.prepare_step_input(request))) + artifacts = tuple(output.close()) + output_closed = True + return artifacts + + +class _Environment: + def __init__(self, *, device: str, world_rank: int, world_size: int) -> None: + self.device = device + self.world_rank = world_rank + self.world_size = world_size + + +def _initialize_environment(device: str) -> _Environment: + """Initialize logging, CUDA placement, and distributed state in the host.""" + 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: + """Console-script entry point.""" + artifacts = run() + for artifact in artifacts: + print(artifact.uri) diff --git a/apps/flashdreams_app/flashdreams_app/contracts.py b/apps/flashdreams_app/flashdreams_app/contracts.py new file mode 100644 index 000000000..f7f1be2a3 --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/contracts.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Data-first provider boundary used by :mod:`flashdreams_app`.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Protocol, runtime_checkable + +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.runtime import InferenceInput, InferenceRuntime +from flashdreams.runtime.types import StepRequest, StepRequirements + + +@dataclass(frozen=True, slots=True) +class AppConfig: + """Provider-specific CLI values normalized by the application host.""" + + options: Mapping[str, Any] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class RuntimeMetadata: + """Presentation facts published by an application runtime.""" + + model_id: str + fps: int | float + output_layout: VideoTensorLayout + video_width: int + video_height: int + + def __post_init__(self) -> None: + if not self.model_id.strip(): + raise ValueError("RuntimeMetadata.model_id must be non-empty.") + if float(self.fps) <= 0: + raise ValueError("RuntimeMetadata.fps must be > 0.") + if self.video_width <= 0 or self.video_height <= 0: + raise ValueError("RuntimeMetadata video dimensions must be > 0.") + + +PipelineCacheInitializer = Callable[[Any, InferenceInput], object] + + +@dataclass(frozen=True, kw_only=True, slots=True) +class PipelineContract: + """Describe the model-specific operation needed to start a pipeline session. + + The host already understands the standard streaming pipeline operations: + ``generate``, ``finalize``, and ``get_num_output_frames``. A provider only + supplies the cache initialization that binds its global conditioning to a + concrete pipeline implementation. + """ + + initialize_cache: PipelineCacheInitializer + + def __post_init__(self) -> None: + if not callable(self.initialize_cache): + raise TypeError("PipelineContract.initialize_cache must be callable.") + + +@dataclass(frozen=True, kw_only=True, slots=True) +class PipelineAppSpec: + """Declarative application definition consumed by the generic host runtime.""" + + pipeline_config: StreamInferencePipelineConfig + contract: PipelineContract + metadata: RuntimeMetadata + initial_input: InferenceInput + total_steps: int + result_metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.pipeline_config, StreamInferencePipelineConfig): + raise TypeError( + "PipelineAppSpec.pipeline_config must be a " + "StreamInferencePipelineConfig." + ) + if not isinstance(self.contract, PipelineContract): + raise TypeError("PipelineAppSpec.contract must be a PipelineContract.") + if not isinstance(self.metadata, RuntimeMetadata): + raise TypeError("PipelineAppSpec.metadata must be RuntimeMetadata.") + if not isinstance(self.initial_input, InferenceInput): + raise TypeError("PipelineAppSpec.initial_input must be InferenceInput.") + if isinstance(self.total_steps, bool) or not isinstance(self.total_steps, int): + raise TypeError("PipelineAppSpec.total_steps must be an integer.") + if self.total_steps <= 0: + raise ValueError("PipelineAppSpec.total_steps must be > 0.") + object.__setattr__( + self, + "result_metadata", + MappingProxyType(dict(self.result_metadata)), + ) + + +def require_pipeline_config( + config: object, *, expected_name: str | None = None +) -> StreamInferencePipelineConfig: + """Validate a pipeline provider result at the application boundary.""" + if not isinstance(config, StreamInferencePipelineConfig): + raise TypeError( + f"Pipeline provider returned {type(config).__name__}, expected " + "StreamInferencePipelineConfig." + ) + if expected_name is not None and config.name != expected_name: + raise ValueError( + f"Preset {expected_name!r} constructed pipeline {config.name!r}; " + "the preset key and pipeline name must match." + ) + return config + + +@runtime_checkable +class AppRuntime(InferenceRuntime, Protocol): + """FlashDreams runtime plus the small surface required by the app host. + + Sessions use the standard :class:`~flashdreams.runtime.InferenceSession` + protocol. Providers add only immutable initial inputs, per-step input + preparation, and presentation metadata. + """ + + metadata: RuntimeMetadata + initial_input: InferenceInput + + def prepare_step_input( + self, request: StepRequest | StepRequirements + ) -> InferenceInput: + """Build model-facing input for one host-owned session step.""" + ... diff --git a/apps/flashdreams_app/flashdreams_app/outputs.py b/apps/flashdreams_app/flashdreams_app/outputs.py new file mode 100644 index 000000000..d6439b45e --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/outputs.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-owned presentation targets.""" + +from __future__ import annotations + +from pathlib import Path + +from flashdreams.infra.postprocess import VideoTensorLayout +from flashdreams.runtime.output import OutputArtifact +from flashdreams.runtime.types import StepResult +from flashdreams.runtime.video_output import Mp4VideoOutputTarget + + +class FileOutput: + """Collect generated chunks and write one MP4 file when the run completes.""" + + def __init__( + self, + *, + path: Path, + fps: int | float, + output_layout: VideoTensorLayout, + enabled: bool = True, + ) -> None: + self._target = Mp4VideoOutputTarget( + output_path=path, + fps=fps, + output_layout=output_layout, + enabled=enabled, + ) + + def open(self) -> None: + """Open the underlying video writer.""" + self._target.open() + + def write(self, result: StepResult) -> None: + """Append a generated chunk.""" + self._target.write(result) + + def close(self) -> tuple[OutputArtifact, ...]: + """Finalize the MP4 and return its artifact metadata.""" + return tuple(self._target.close()) diff --git a/apps/flashdreams_app/flashdreams_app/runtime.py b/apps/flashdreams_app/flashdreams_app/runtime.py new file mode 100644 index 000000000..f5d9546a3 --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/runtime.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-owned runtime and session for streaming pipeline applications.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from flashdreams.infra.config import derive_config +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.runtime import ( + InferenceInput, + InferenceSession, + StepRequest, + StepRequirements, + StepResult, +) + +from .contracts import PipelineAppSpec, RuntimeMetadata + + +class PipelineAppRuntime: + """Instantiate and own the reusable pipeline behind an application spec.""" + + def __init__( + self, + *, + spec: PipelineAppSpec, + device: str, + compile: bool | None = None, + cuda_graph: bool | None = None, + ) -> None: + pipeline_config = _with_execution_options( + spec.pipeline_config, + compile=compile, + cuda_graph=cuda_graph, + ) + self.pipeline = pipeline_config.setup().to(device).eval() + self.metadata = spec.metadata + self.initial_input = spec.initial_input + self._spec = spec + self._closed = False + + def prepare_step_input( + self, request: StepRequest | StepRequirements + ) -> InferenceInput: + """Return an empty step payload for prompt-conditioned pipelines.""" + del request + return InferenceInput() + + def start_session(self, inputs: InferenceInput) -> "PipelineAppSession": + """Create a finite session with cache state isolated from other sessions.""" + if self._closed: + raise RuntimeError("Pipeline application runtime is closed.") + return PipelineAppSession( + pipeline=self.pipeline, + inputs=inputs, + spec=self._spec, + ) + + def peek_input_fps(self) -> float: + """Return the host clock rate used for realtime presentation.""" + return float(self.metadata.fps) + + def peek_steady_output_num_frames(self) -> int: + """Return the steady-state output chunk size for presentation queues.""" + return int(self.pipeline.get_num_output_frames(1)) + + def close(self) -> None: + """Release the shared pipeline and accelerator allocator state.""" + if self._closed: + return + self._closed = True + close = getattr(self.pipeline, "close", None) + if callable(close): + close() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +class PipelineAppSession(InferenceSession): + """Host-owned finite autoregressive session for a pipeline app spec.""" + + def __init__( + self, + *, + pipeline: Any, + inputs: InferenceInput, + spec: PipelineAppSpec, + ) -> None: + self._pipeline = pipeline + self._cache: object | None = spec.contract.initialize_cache(pipeline, inputs) + self._metadata = spec.metadata + self._result_metadata = spec.result_metadata + self._total_steps = spec.total_steps + self._step_index = 0 + self._closed = False + + def next_step_request(self) -> StepRequest | None: + """Return the next finite rollout request, or ``None`` when complete.""" + if self._closed or self._step_index >= self._total_steps: + return None + return StepRequest(step_index=self._step_index) + + def step(self, inputs: InferenceInput) -> StepResult: + """Generate and finalize one autoregressive video chunk.""" + del inputs + if self._closed: + raise RuntimeError("Pipeline application session is closed.") + if self._step_index >= self._total_steps: + raise RuntimeError("Pipeline application session is complete.") + if self._cache is None: + raise RuntimeError("Pipeline application session has no active cache.") + + index = self._step_index + video = self._pipeline.generate( + autoregressive_index=index, + cache=self._cache, + ) + metrics = _metrics( + self._pipeline.finalize( + autoregressive_index=index, + cache=self._cache, + ) + ) + self._step_index += 1 + if not isinstance(video, torch.Tensor): + raise TypeError( + "Pipeline generate() must return a torch.Tensor, got " + f"{type(video).__name__}." + ) + return StepResult.from_video_chunk( + step_index=index, + video_chunk=video.detach(), + layout=self._metadata.output_layout, + metrics=metrics, + metadata=self._result_metadata, + ) + + def reset(self, inputs: InferenceInput | None = None) -> None: + """Reject reset because finite sessions use isolated cache state.""" + del inputs + raise RuntimeError("Create a new session instead of resetting this one.") + + def close(self) -> None: + """Release session-local cache state.""" + if self._closed: + return + self._closed = True + cache = self._cache + self._cache = None + close = getattr(cache, "close", None) + if callable(close): + close() + + +def _with_execution_options( + pipeline: StreamInferencePipelineConfig, + *, + compile: bool | None, + cuda_graph: bool | None, +) -> StreamInferencePipelineConfig: + transformer: dict[str, object] = {} + if compile is not None: + transformer["compile_network"] = compile + if cuda_graph is not None: + transformer["use_cuda_graph"] = cuda_graph + if not transformer: + return pipeline + return derive_config( + pipeline, + diffusion_model={"transformer": transformer}, + ) + + +def _metrics(value: object) -> Mapping[str, float | int]: + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError( + "Pipeline finalize() must return a metrics mapping or None, got " + f"{type(value).__name__}." + ) + metrics = dict(value) + invalid = tuple( + key for key, metric in metrics.items() if not isinstance(metric, (int, float)) + ) + if invalid: + raise TypeError(f"Pipeline metrics must be numeric; invalid keys: {invalid}.") + return metrics + + +__all__ = ["PipelineAppRuntime", "PipelineAppSession"] diff --git a/apps/flashdreams_app/flashdreams_app/webrtc.py b/apps/flashdreams_app/flashdreams_app/webrtc.py new file mode 100644 index 000000000..32f6d8a4f --- /dev/null +++ b/apps/flashdreams_app/flashdreams_app/webrtc.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Host-owned WebRTC presentation for application runtimes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from flashdreams.runtime import InferenceConfig, InferenceInput +from flashdreams.runtime.demo import ( + DemoSpec, + 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 .contracts import AppRuntime + + +@dataclass(frozen=True, slots=True) +class WebRTCOptions: + """Presentation settings owned by ``flashdreams-app``.""" + + host: str + port: int + warmup_chunks: int + warmup_timeout_s: float + client_liveness_timeout_s: float + device: str + encoder_backend: str + encoder_bitrate_bps: int + encoder_gop: int + + +@dataclass(frozen=True, slots=True) +class _WebRTCRuntimeConfig: + video_width: int + video_height: int + warmup_chunks: int + warmup_timeout_s: float + device: str + encoder_backend: str + encoder_bitrate_bps: int + encoder_gop: int + + +class _InputProvider: + capabilities = ProviderCapabilities( + supports_realtime_clock=True, + deterministic_given_inputs=True, + ) + + def __init__(self, runtime: AppRuntime) -> None: + self._runtime = runtime + + def prepare_initial_input(self) -> InferenceInput: + return self._runtime.initial_input + + def prepare_step( + self, *, request: StepRequirements, user_window: UserInputWindow + ) -> PreparedStep: + del user_window + return PreparedStep(inference_input=self._runtime.prepare_step_input(request)) + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + pass + + +def serve_webrtc( + *, runtime: AppRuntime, options: WebRTCOptions, world_rank: int +) -> object: + """Serve an application runtime through the shared WebRTC stack.""" + metadata = runtime.metadata + output = WebRTCOutputSpec( + host=options.host, + port=options.port, + fps=int(metadata.fps), + video_width=metadata.video_width, + video_height=metadata.video_height, + warmup_chunks=options.warmup_chunks, + warmup_timeout_s=options.warmup_timeout_s, + client_liveness_timeout_s=options.client_liveness_timeout_s, + preload_name=metadata.model_id, + ) + spec = DemoSpec( + model_id=metadata.model_id, + input_mode="webrtc", + output=output, + config=InferenceConfig(model_id=metadata.model_id, device=options.device), + ) + scenario = PreparedScenario(initial_inputs=runtime.initial_input) + + def create_model_input_provider( + spec: DemoSpec, + scenario: PreparedScenario, + ) -> _InputProvider: + del spec, scenario + return _InputProvider(runtime) + + manager = BaseWebRTCSessionManager( + runtime=runtime, + runtime_config=_WebRTCRuntimeConfig( + video_width=metadata.video_width, + video_height=metadata.video_height, + warmup_chunks=options.warmup_chunks, + warmup_timeout_s=options.warmup_timeout_s, + device=options.device, + encoder_backend=options.encoder_backend, + encoder_bitrate_bps=options.encoder_bitrate_bps, + encoder_gop=options.encoder_gop, + ), + fps=int(metadata.fps), + identity=metadata.model_id, + shared_host=RuntimeHost(runtime), + shared_spec=spec, + shared_scenario=scenario, + shared_model_input_provider_factory=create_model_input_provider, + client_liveness_timeout_s=options.client_liveness_timeout_s, + ) + return serve_webrtc_demo( + output=output, + model_id=metadata.model_id, + session_manager=manager, + app_resources=WebRTCAppResources(preload_name=metadata.model_id), + world_rank=world_rank, + ) diff --git a/apps/flashdreams_app/pyproject.toml b/apps/flashdreams_app/pyproject.toml new file mode 100644 index 000000000..511814b94 --- /dev/null +++ b/apps/flashdreams_app/pyproject.toml @@ -0,0 +1,22 @@ +# 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-app" +version = "0.1.0" +description = "Generic host for FlashDreams application runtimes" +requires-python = ">=3.10" +dependencies = ["flashdreams[serving]"] + +[project.scripts] +flashdreams-app = "flashdreams_app.cli:main" + +[tool.uv.sources] +flashdreams = { workspace = true } + +[tool.setuptools.packages.find] +where = ["."] diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py new file mode 100644 index 000000000..55407cecb --- /dev/null +++ b/apps/flashdreams_app/tests/test_cli.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import ModuleType, SimpleNamespace +from typing import Any, cast + +import pytest +import torch +from flashdreams_app import ( + PipelineAppRuntime, + PipelineAppSpec, + PipelineContract, + RuntimeMetadata, + cli, +) + +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.runtime import InferenceInput, StepResult + +pytestmark = pytest.mark.ci_cpu + + +def test_host_drives_runtime_api_and_owns_file_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + class Cache: + def close(self) -> None: + calls.append("cache.close") + + class Pipeline: + def __init__(self, config: object) -> None: + del config + calls.append("pipeline.init") + + def to(self, device: str) -> "Pipeline": + assert device == "cpu" + calls.append("pipeline.to") + return self + + def eval(self) -> "Pipeline": + calls.append("pipeline.eval") + return self + + def generate(self, *, autoregressive_index: int, cache: object) -> torch.Tensor: + assert autoregressive_index == 0 + assert isinstance(cache, Cache) + calls.append("pipeline.generate") + return torch.zeros((1, 3, 2, 2)) + + def finalize( + self, *, autoregressive_index: int, cache: object + ) -> dict[str, float]: + assert autoregressive_index == 0 + assert isinstance(cache, Cache) + calls.append("pipeline.finalize") + return {"step_ms": 1.0} + + def close(self) -> None: + calls.append("pipeline.close") + + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, Pipeline), + name="fake", + diffusion_model=cast(Any, None), + ) + + def initialize_cache(pipeline: object, inputs: InferenceInput) -> object: + assert isinstance(pipeline, Pipeline) + assert inputs.global_conditioning["prompt"] == "test" + calls.append("contract.initialize_cache") + return Cache() + + provider = ModuleType("fake_app") + setattr( + provider, + "create_app", + lambda config: PipelineAppSpec( + pipeline_config=pipeline_config, + contract=PipelineContract(initialize_cache=initialize_cache), + metadata=RuntimeMetadata( + model_id="fake", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + ), + initial_input=InferenceInput(global_conditioning={"prompt": "test"}), + total_steps=1, + ), + ) + monkeypatch.setattr(cli, "load_provider", lambda _: provider) + + class Output: + def __init__(self, **_: object) -> None: + calls.append("output.init") + + def open(self) -> None: + calls.append("output.open") + + def write(self, result: StepResult) -> None: + calls.append("output.write") + + def close(self) -> tuple[object, ...]: + calls.append("output.close") + return () + + monkeypatch.setattr(cli, "FileOutput", Output) + cli.run(["fake-app", "mp4", "--device", "cpu", "--output", "result.mp4"]) + assert calls == [ + "pipeline.init", + "pipeline.to", + "pipeline.eval", + "output.init", + "output.open", + "contract.initialize_cache", + "pipeline.generate", + "pipeline.finalize", + "output.write", + "output.close", + "cache.close", + "pipeline.close", + ] + + +def test_host_exposes_only_supported_output_modes() -> None: + mode_action = next( + action for action in cli.build_parser()._actions if action.dest == "mode" + ) + assert mode_action.choices == ("mp4", "webrtc") + + +def test_host_owns_execution_options() -> None: + destinations = {action.dest for action in cli.build_parser()._actions} + assert {"compile", "cuda_graph"} <= destinations + + +def test_host_applies_execution_options_without_mutating_provider_spec() -> None: + configured_pipeline: StreamInferencePipelineConfig | None = None + + class Pipeline: + def __init__(self, config: StreamInferencePipelineConfig) -> None: + nonlocal configured_pipeline + configured_pipeline = config + + def to(self, device: str) -> "Pipeline": + assert device == "cpu" + return self + + def eval(self) -> "Pipeline": + return self + + transformer = SimpleNamespace(compile_network=False, use_cuda_graph=False) + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, Pipeline), + name="fake", + diffusion_model=cast(Any, SimpleNamespace(transformer=transformer)), + ) + spec = PipelineAppSpec( + pipeline_config=pipeline_config, + contract=PipelineContract(initialize_cache=lambda pipeline, inputs: object()), + metadata=RuntimeMetadata( + model_id="fake", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + ), + initial_input=InferenceInput(), + total_steps=1, + ) + + runtime = PipelineAppRuntime( + spec=spec, + device="cpu", + compile=True, + cuda_graph=True, + ) + try: + assert configured_pipeline is not None + configured_transformer = configured_pipeline.diffusion_model.transformer + assert getattr(configured_transformer, "compile_network") is True + assert getattr(configured_transformer, "use_cuda_graph") is True + assert transformer.compile_network is False + assert transformer.use_cuda_graph is False + finally: + runtime.close() diff --git a/apps/flashdreams_app/tests/test_webrtc.py b/apps/flashdreams_app/tests/test_webrtc.py new file mode 100644 index 000000000..008e68d60 --- /dev/null +++ b/apps/flashdreams_app/tests/test_webrtc.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +from flashdreams_app import RuntimeMetadata, webrtc + +from flashdreams.runtime import InferenceInput, StepRequest, StepResult + +pytestmark = pytest.mark.ci_cpu + + +class _Session: + def next_step_request(self) -> StepRequest | None: + return None + + def step(self, inputs: InferenceInput) -> StepResult: + raise AssertionError("Server construction must not step a session.") + + def reset(self, inputs: InferenceInput | None = None) -> None: + del inputs + + def close(self) -> None: + pass + + +class _Runtime: + metadata = RuntimeMetadata( + model_id="fake-app", + fps=16, + output_layout="tchw", + video_width=96, + video_height=64, + ) + initial_input = InferenceInput() + + def prepare_step_input(self, request: object) -> InferenceInput: + del request + return InferenceInput() + + def start_session(self, inputs: InferenceInput) -> _Session: + del inputs + return _Session() + + def close(self) -> None: + pass + + +def test_host_constructs_webrtc_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, + options=webrtc.WebRTCOptions( + host="127.0.0.1", + port=8080, + warmup_chunks=0, + warmup_timeout_s=30.0, + client_liveness_timeout_s=30.0, + device="cpu", + encoder_backend="default", + encoder_bitrate_bps=1_000_000, + encoder_gop=16, + ), + world_rank=0, + ) + + assert result == "served" + assert captured["model_id"] == "fake-app" + assert captured["world_rank"] == 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._shared_scenario is not None + assert session_manager._shared_scenario.initial_inputs is runtime.initial_input + assert callable(session_manager._shared_model_input_provider_factory) diff --git a/apps/t2v_app/README.md b/apps/t2v_app/README.md new file mode 100644 index 000000000..c17ce47a1 --- /dev/null +++ b/apps/t2v_app/README.md @@ -0,0 +1,50 @@ +# T2V App Provider + +`t2v-app` is a model provider for the generic `flashdreams-app` host. It +returns a declarative `PipelineAppSpec`; it does not implement a runtime or +session and does not own setup, stepping, finalization, cleanup, MP4 writing, +or WebRTC. + +The provider loads a YAML preset catalog through +`flashdreams.core.pipeline_presets` and asks the selected pipeline provider to +construct a FlashDreams `StreamInferencePipelineConfig`. It does not use the +`flashdreams.runner_configs` registry. The packaged catalog is +[`t2v_app/pipeline_presets.yaml`](t2v_app/pipeline_presets.yaml); pass +`--preset-config` to use another catalog. + +```bash +uv run flashdreams-app 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-app t2v-app webrtc \ + --preset-id self-forcing-wan2.1-t2v-1.3b \ + --prompt "A neon-lit city at night" + +uv run flashdreams-app t2v-app mp4 \ + --preset-config /path/to/presets.yaml \ + --preset-id my-t2v-preset \ + --prompt "A waterfall" \ + --output outputs/waterfall.mp4 +``` + +Every YAML preset must specify `provider`, all six runtime/presentation fields, +and the provider-owned `pipeline` options. FlashDreams' +`ObjectGraphPipelineProvider` supports these 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 provider class (or provider +instance) implementing `flashdreams.core.pipeline_presets.PipelineProvider` +and reference it from `provider`. Preset YAML is trusted configuration because +provider and object-graph references import Python objects. + +At the provider boundary, T2V contributes only its preset selection, +conditioning values, presentation metadata, and a cache initializer that maps +the prompt and pixel dimensions to the selected pipeline. `flashdreams-app` +constructs and drives the resulting pipeline. diff --git a/apps/t2v_app/pyproject.toml b/apps/t2v_app/pyproject.toml new file mode 100644 index 000000000..0e0b61081 --- /dev/null +++ b/apps/t2v_app/pyproject.toml @@ -0,0 +1,24 @@ +# 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 pipeline provider for flashdreams-app" +readme = "README.md" +requires-python = ">=3.10" +dependencies = ["flashdreams", "flashdreams-app"] + +[tool.uv.sources] +flashdreams = { workspace = true } +flashdreams-app = { workspace = true } + +[tool.setuptools.packages.find] +where = ["."] + +[tool.setuptools.package-data] +t2v_app = ["pipeline_presets.yaml"] diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py new file mode 100644 index 000000000..d4aff7bd6 --- /dev/null +++ b/apps/t2v_app/t2v_app/__init__.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text-to-video pipeline application for ``flashdreams-app``.""" + +from flashdreams.core.pipeline_presets import PipelineProvider + +from .provider import add_arguments, create_app + +__all__ = ["PipelineProvider", "add_arguments", "create_app"] 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..60b5e30f3 --- /dev/null +++ b/apps/t2v_app/t2v_app/presets.py @@ -0,0 +1,159 @@ +# 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 + +_RUNTIME_FIELDS = { + "prompt", + "total_blocks", + "pixel_height", + "pixel_width", + "fps", + "output_layout", +} +_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.""" + + total_blocks: int + """Number of autoregressive chunks in one finite session.""" + + 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.""" + + +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_exact_fields(runtime, expected=_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"), + total_blocks=_positive_int( + runtime["total_blocks"], path=f"{path}.total_blocks" + ), + 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), + ) + + +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() + + +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/provider.py b/apps/t2v_app/t2v_app/provider.py new file mode 100644 index 000000000..4fb0f3f3e --- /dev/null +++ b/apps/t2v_app/t2v_app/provider.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Text-to-video application definition for the FlashDreams app host.""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +from flashdreams_app import ( + AppConfig, + PipelineAppSpec, + PipelineContract, + RuntimeMetadata, + require_pipeline_config, +) + +from flashdreams.core.pipeline_presets import load_pipeline_provider +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.runtime import InferenceInput + +from .presets import RuntimePresetOptions, load_preset_catalog + +FIELD_PROMPT = "prompt" +FIELD_TOTAL_BLOCKS = "total_blocks" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" +FIELD_FPS = "fps" + + +def add_arguments(parser: argparse.ArgumentParser) -> None: + """Register T2V conditioning and rollout options on the host 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) + + +def create_app(config: AppConfig) -> PipelineAppSpec: + """Describe a T2V pipeline application without constructing its runtime.""" + options = config.options + catalog = load_preset_catalog(_optional_path(options.get("preset_config"))) + preset_id, preset = catalog.resolve(_optional_string(options.get("preset_id"))) + provider = load_pipeline_provider(preset.provider) + pipeline_config = require_pipeline_config( + provider.create_pipeline_config( + preset_id=preset_id, + options=preset.pipeline, + ), + expected_name=preset_id, + ) + scenario = _scenario(options, preset.runtime) + return PipelineAppSpec( + pipeline_config=pipeline_config, + contract=PipelineContract(initialize_cache=_initialize_cache), + metadata=RuntimeMetadata( + 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 + ), + ), + initial_input=InferenceInput(global_conditioning=scenario), + total_steps=_required_int( + scenario[FIELD_TOTAL_BLOCKS], name=FIELD_TOTAL_BLOCKS + ), + result_metadata={FIELD_PROMPT: scenario[FIELD_PROMPT]}, + ) + + +def _initialize_cache(pipeline: Any, inputs: InferenceInput) -> object: + """Bind T2V prompt and dimensions to a new pipeline cache.""" + scenario = _scenario_from_inputs(inputs) + decoder = pipeline.decoder + if not isinstance(decoder, StreamingVideoDecoder): + raise TypeError("T2V pipelines require a StreamingVideoDecoder.") + ratio = decoder.spatial_compression_ratio + pixel_height = _required_int(scenario[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT) + pixel_width = _required_int(scenario[FIELD_PIXEL_WIDTH], name=FIELD_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})." + ) + return pipeline.initialize_cache( + text=[str(scenario[FIELD_PROMPT])], + image=None, + height=pixel_height // ratio, + width=pixel_width // ratio, + ) + + +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_TOTAL_BLOCKS: _option_or_default( + options, FIELD_TOTAL_BLOCKS, defaults.total_blocks + ), + 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_TOTAL_BLOCKS, + 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 _scenario_from_inputs(inputs: InferenceInput) -> Mapping[str, object]: + source = inputs.global_conditioning + required = ( + FIELD_PROMPT, + FIELD_TOTAL_BLOCKS, + FIELD_PIXEL_HEIGHT, + FIELD_PIXEL_WIDTH, + FIELD_FPS, + ) + missing = tuple(name for name in required if name not in source) + if missing: + raise ValueError(f"Missing T2V global conditioning fields: {missing}.") + scenario = dict(source) + scenario[FIELD_PROMPT] = _resolve_prompt(scenario[FIELD_PROMPT]) + for name in ( + FIELD_TOTAL_BLOCKS, + 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 _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__ = ["add_arguments", "create_app"] 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..a069bffe4 --- /dev/null +++ b/apps/t2v_app/tests/test_pipeline_provider.py @@ -0,0 +1,88 @@ +# 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 flashdreams_app import require_pipeline_config +from t2v_app.presets import load_preset_catalog + +from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict +from flashdreams.core.pipeline_presets import load_pipeline_provider +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) + + provider = load_pipeline_provider(preset.provider) + config = require_pipeline_config( + provider.create_pipeline_config( + preset_id=preset_id, + options=preset.pipeline, + ), + expected_name=preset_id, + ) + + 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_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_provider.py b/apps/t2v_app/tests/test_provider.py new file mode 100644 index 000000000..73c42dc3c --- /dev/null +++ b/apps/t2v_app/tests/test_provider.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the T2V application provider boundary.""" + +from __future__ import annotations + +import argparse +from typing import Any, cast + +import pytest +from flashdreams_app import AppConfig, PipelineAppSpec +from t2v_app import provider +from t2v_app.presets import ( + PipelinePreset, + PresetCatalog, + RuntimePresetOptions, +) + +from flashdreams.infra.pipeline import StreamInferencePipelineConfig + +pytestmark = pytest.mark.ci_cpu + + +def test_t2v_provider_registers_model_options() -> None: + parser = argparse.ArgumentParser() + provider.add_arguments(parser) + args = parser.parse_args(["--prompt", "A waterfall"]) + assert args.prompt == "A waterfall" + assert args.preset_id is None + assert args.preset_config is None + assert not hasattr(args, "backend") + assert not hasattr(args, "compile") + + +def test_create_app_returns_data_without_constructing_pipeline( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pipeline_constructed = False + + class Pipeline: + def __init__(self, _: object) -> None: + nonlocal pipeline_constructed + pipeline_constructed = True + + preset_id = "test-t2v" + defaults = RuntimePresetOptions( + prompt="default prompt", + total_blocks=2, + pixel_height=64, + pixel_width=96, + fps=12, + output_layout="tchw", + ) + preset = PipelinePreset( + provider="tests:provider", + runtime=defaults, + pipeline={}, + ) + catalog = PresetCatalog( + default_preset_id=preset_id, + presets={preset_id: preset}, + ) + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, Pipeline), + name=preset_id, + diffusion_model=cast(Any, None), + ) + + class Provider: + def create_pipeline_config( + self, *, preset_id: str, options: object + ) -> StreamInferencePipelineConfig: + assert preset_id == "test-t2v" + assert options == {} + return pipeline_config + + monkeypatch.setattr(provider, "load_preset_catalog", lambda _: catalog) + monkeypatch.setattr(provider, "load_pipeline_provider", lambda _: Provider()) + + created = provider.create_app( + AppConfig( + options={ + "preset_config": None, + "preset_id": None, + "prompt": "A waterfall", + "total_blocks": None, + "pixel_height": None, + "pixel_width": None, + "fps": None, + } + ) + ) + + assert isinstance(created, PipelineAppSpec) + assert created.pipeline_config is pipeline_config + assert created.initial_input.global_conditioning["prompt"] == "A waterfall" + assert created.metadata.video_width == 96 + assert created.metadata.fps == 12 + assert created.total_steps == 2 + assert not pipeline_constructed 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..db22b87a7 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -680,6 +680,9 @@ 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, @@ -713,6 +716,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 +764,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 +905,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) @@ -1886,11 +1897,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 +1919,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 +1936,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/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..803133510 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,8 @@ no-build-isolation-package = ["transformer-engine-torch"] extraPaths = [ "flashdreams", "apps", + "apps/flashdreams_app", + "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", "integrations/causal_forcing", @@ -70,6 +72,8 @@ python-version = "3.10" extra-paths = [ "flashdreams", "apps", + "apps/flashdreams_app", + "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", "integrations/causal_forcing", diff --git a/uv.lock b/uv.lock index 7f4aba302..e275dc9f7 100644 --- a/uv.lock +++ b/uv.lock @@ -20,6 +20,7 @@ conflicts = [[ [manifest] members = [ "flashdreams", + "flashdreams-app", "flashdreams-causal-forcing", "flashdreams-cosmos-predict2", "flashdreams-fastvideo-causal-wan22", @@ -33,6 +34,7 @@ members = [ "flashdreams-wan21", "flashdreams-wan22", "ludus-renderer", + "t2v-app", ] overrides = [ { name = "numpy", specifier = ">=1.24,<2.5" }, @@ -1107,6 +1109,17 @@ cuda13 = [ { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, ] +[[package]] +name = "flashdreams-app" +version = "0.1.0" +source = { editable = "apps/flashdreams_app" } +dependencies = [ + { name = "flashdreams", extra = ["serving"] }, +] + +[package.metadata] +requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] + [[package]] name = "flashdreams-causal-forcing" 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-app" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-app", editable = "apps/flashdreams_app" }, +] + [[package]] name = "tokenizers" version = "0.22.2" From 3e03ff1f14ebdd607067184487d709823635fa92 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 02:37:19 +0000 Subject: [PATCH 02/10] Separate FlashDreams app presentation paths Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 90 +++++++++++++++++++ apps/flashdreams_app/flashdreams_app/cli.py | 95 ++++++++++++++++----- apps/flashdreams_app/tests/test_cli.py | 58 +++++++++++++ 3 files changed, 221 insertions(+), 22 deletions(-) create mode 100644 PR_DESCRIPTION.md diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md new file mode 100644 index 000000000..775e3215a --- /dev/null +++ b/PR_DESCRIPTION.md @@ -0,0 +1,90 @@ +## Summary + +Add a model-neutral `flashdreams-app` entrypoint and a declarative `t2v-app` +provider for running FlashDreams text-to-video pipelines as MP4 jobs or WebRTC +services. + +The host owns process initialization, pipeline construction, execution options, +the runtime/session lifecycle, autoregressive stepping, finalization, cleanup, +and presentation. The T2V package only selects a pipeline preset and supplies +the model-specific conditioning and cache-initialization contract. + +## Entrypoint examples + +Generate an MP4 with the packaged default preset: + +```bash +uv run flashdreams-app t2v-app mp4 \ + --prompt "A waterfall" \ + --output o.mp4 +``` + +Serve the same application through WebRTC: + +```bash +uv run flashdreams-app t2v-app webrtc \ + --prompt "A waterfall" +``` + +Select a packaged preset explicitly: + +```bash +uv run flashdreams-app t2v-app mp4 \ + --preset-id self-forcing-wan2.1-t2v-1.3b \ + --prompt "A neon-lit city at night" \ + --output outputs/city.mp4 +``` + +Load a custom YAML preset catalog: + +```bash +uv run flashdreams-app t2v-app webrtc \ + --preset-config /path/to/pipeline-presets.yaml \ + --preset-id my-t2v-preset \ + --prompt "A waterfall at sunset" +``` + +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`. + +## Architecture + +- Add the `flashdreams-app` workspace package and console entrypoint. +- Define a minimal provider boundary: + `create_app(AppConfig) -> PipelineAppSpec`. +- Add the host-owned `PipelineAppRuntime` and `PipelineAppSession`. +- Keep pipeline setup, `generate`/`finalize`, step tracking, cache release, and + runtime closure in the host. +- Add host-owned MP4 and WebRTC presentation paths without a runtime adapter. +- Keep provider-specific CLI arguments optional through `add_arguments(parser)`. +- Move `--compile` and `--cuda-graph` execution overrides to the host. + +## T2V provider and presets + +- Add the `t2v-app` workspace package. +- Describe T2V through a `PipelineAppSpec` rather than implementing another + runtime or session. +- Add packaged YAML presets for causal-forcing and self-forcing WAN pipelines. +- Resolve pipeline providers directly from the YAML catalog without depending + on the runner registry. +- Support trusted declarative `_target`, `_ref`, and `_tuple` nodes for pipeline + object graphs. + +## 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 the causal-forcing and self-forcing configurations. +- Let the shared asynchronous demo and WebRTC manager consume runtime objects + directly when no model adapter is needed. + +## Validation + +- `57` affected CPU tests pass. +- Full `pre-commit run -a` passes, including Ruff formatting, lockfile + validation, and `ty` type checking. +- `flashdreams-app t2v-app mp4 --help` resolves the installed provider and + composes host and T2V arguments correctly. +- GPU model generation was not run as part of this change. diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py index 8058e5414..2d3e45aa9 100644 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -21,7 +21,7 @@ ) from flashdreams.runtime.output import OutputArtifact -from .contracts import AppConfig, PipelineAppSpec +from .contracts import AppConfig, AppRuntime, PipelineAppSpec from .outputs import FileOutput from .runtime import PipelineAppRuntime from .webrtc import WebRTCOptions, serve_webrtc @@ -83,7 +83,14 @@ def load_provider(distribution_name: str) -> ModuleType: def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: - """Run one provider session and return artifacts produced by the host.""" + """Dispatch one provider session to its selected presentation path. + + Args: + argv: Command-line arguments; ``None`` reads the process arguments. + + Returns: + Artifacts produced by the selected path. WebRTC returns an empty tuple. + """ probe = argparse.ArgumentParser(add_help=False) probe.add_argument("provider", nargs="?") probe.add_argument("mode", nargs="?") @@ -122,30 +129,74 @@ def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: cuda_graph=args.cuda_graph, ) if args.mode == "webrtc": - try: - serve_webrtc( - runtime=runtime, - options=WebRTCOptions( - host=args.host, - port=args.port, - warmup_chunks=args.warmup_chunks, - warmup_timeout_s=args.warmup_timeout_s, - client_liveness_timeout_s=args.client_liveness_timeout_s, - device=environment.device, - encoder_backend=args.encoder_backend, - encoder_bitrate_bps=args.encoder_bitrate_bps, - encoder_gop=args.encoder_gop or int(runtime.metadata.fps), - ), - world_rank=environment.world_rank, - ) - finally: - runtime.close() - return () + return _run_webrtc(runtime=runtime, args=args, environment=environment) + if args.mode == "mp4": + return _run_mp4( + runtime=runtime, + output_path=args.output, + environment=environment, + ) + raise AssertionError(f"Unsupported presentation mode: {args.mode!r}.") + + +def _run_webrtc( + *, + runtime: AppRuntime, + args: argparse.Namespace, + environment: "_Environment", +) -> tuple[OutputArtifact, ...]: + """Run the WebRTC serving path and close its runtime. + + Args: + runtime: Initialized application runtime. + args: Parsed host and WebRTC arguments. + environment: Initialized process and distributed environment. + + Returns: + An empty tuple because WebRTC does not create file artifacts. + """ + try: + serve_webrtc( + runtime=runtime, + options=WebRTCOptions( + host=args.host, + port=args.port, + warmup_chunks=args.warmup_chunks, + warmup_timeout_s=args.warmup_timeout_s, + client_liveness_timeout_s=args.client_liveness_timeout_s, + device=environment.device, + encoder_backend=args.encoder_backend, + encoder_bitrate_bps=args.encoder_bitrate_bps, + encoder_gop=args.encoder_gop or int(runtime.metadata.fps), + ), + world_rank=environment.world_rank, + ) + finally: + runtime.close() + return () + + +def _run_mp4( + *, + runtime: AppRuntime, + output_path: Path, + environment: "_Environment", +) -> tuple[OutputArtifact, ...]: + """Run the finite MP4 generation path and close all owned resources. + + Args: + runtime: Initialized application runtime. + output_path: Destination for the generated MP4. + environment: Initialized process and distributed environment. + + Returns: + Artifacts emitted by the file output target. + """ with ExitStack() as resources: resources.callback(runtime.close) output = FileOutput( - path=args.output, + path=output_path, fps=runtime.metadata.fps, output_layout=runtime.metadata.output_layout, enabled=environment.world_rank == 0, diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py index 55407cecb..261fbac22 100644 --- a/apps/flashdreams_app/tests/test_cli.py +++ b/apps/flashdreams_app/tests/test_cli.py @@ -3,6 +3,7 @@ from __future__ import annotations +import argparse from types import ModuleType, SimpleNamespace from typing import Any, cast @@ -133,6 +134,63 @@ def test_host_exposes_only_supported_output_modes() -> None: assert mode_action.choices == ("mp4", "webrtc") +def test_webrtc_path_owns_serving_options_and_runtime_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + captured: dict[str, object] = {} + + class Runtime: + metadata = RuntimeMetadata( + model_id="fake", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + ) + initial_input = InferenceInput() + + def prepare_step_input(self, request: object) -> InferenceInput: + del request + return InferenceInput() + + def start_session(self, inputs: InferenceInput) -> Any: + del inputs + raise AssertionError("The WebRTC path must not start an MP4 session.") + + def close(self) -> None: + calls.append("runtime.close") + + def serve(**kwargs: object) -> None: + calls.append("serve_webrtc") + captured.update(kwargs) + + monkeypatch.setattr(cli, "serve_webrtc", serve) + result = cli._run_webrtc( + runtime=Runtime(), + args=argparse.Namespace( + host="127.0.0.1", + port=9000, + warmup_chunks=2, + warmup_timeout_s=30.0, + client_liveness_timeout_s=10.0, + encoder_backend="default", + encoder_bitrate_bps=1_000_000, + encoder_gop=None, + ), + environment=cli._Environment(device="cpu", world_rank=0, world_size=1), + ) + + assert result == () + assert calls == ["serve_webrtc", "runtime.close"] + assert captured["world_rank"] == 0 + options = captured["options"] + assert isinstance(options, cli.WebRTCOptions) + assert options.host == "127.0.0.1" + assert options.port == 9000 + assert options.encoder_gop == 24 + + def test_host_owns_execution_options() -> None: destinations = {action.dest for action in cli.build_parser()._actions} assert {"compile", "cuda_graph"} <= destinations From b548c6ec5cf01526f0cf49803c076fd7cd9a022c Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 02:50:10 +0000 Subject: [PATCH 03/10] Keep pipeline setup encapsulated Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 13 ++-- apps/flashdreams_app/README.md | 12 ++-- apps/flashdreams_app/flashdreams_app/cli.py | 16 ++--- .../flashdreams_app/runtime.py | 30 +------- apps/flashdreams_app/tests/test_cli.py | 70 +++---------------- apps/t2v_app/t2v_app/__init__.py | 4 +- apps/t2v_app/t2v_app/provider.py | 4 +- apps/t2v_app/tests/test_provider.py | 4 +- 8 files changed, 35 insertions(+), 118 deletions(-) diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md index 775e3215a..6808b6647 100644 --- a/PR_DESCRIPTION.md +++ b/PR_DESCRIPTION.md @@ -4,10 +4,10 @@ Add a model-neutral `flashdreams-app` entrypoint and a declarative `t2v-app` provider for running FlashDreams text-to-video pipelines as MP4 jobs or WebRTC services. -The host owns process initialization, pipeline construction, execution options, -the runtime/session lifecycle, autoregressive stepping, finalization, cleanup, -and presentation. The T2V package only selects a pipeline preset and supplies -the model-specific conditioning and cache-initialization contract. +The host owns process initialization, pipeline construction, the +runtime/session lifecycle, autoregressive stepping, finalization, cleanup, and +presentation. The T2V package only selects a pipeline preset and supplies the +model-specific conditioning and cache-initialization contract. ## Entrypoint examples @@ -52,13 +52,14 @@ When `--preset-id` is omitted, `t2v-app` uses the catalog's - Add the `flashdreams-app` workspace package and console entrypoint. - Define a minimal provider boundary: - `create_app(AppConfig) -> PipelineAppSpec`. + `create_app_spec(AppConfig) -> PipelineAppSpec`. - Add the host-owned `PipelineAppRuntime` and `PipelineAppSession`. - Keep pipeline setup, `generate`/`finalize`, step tracking, cache release, and runtime closure in the host. - Add host-owned MP4 and WebRTC presentation paths without a runtime adapter. - Keep provider-specific CLI arguments optional through `add_arguments(parser)`. -- Move `--compile` and `--cuda-graph` execution overrides to the host. +- Keep pipeline-specific execution behavior encapsulated by the selected + pipeline config and its `setup()` implementation. ## T2V provider and presets diff --git a/apps/flashdreams_app/README.md b/apps/flashdreams_app/README.md index b1ca3d872..c8b697291 100644 --- a/apps/flashdreams_app/README.md +++ b/apps/flashdreams_app/README.md @@ -14,13 +14,15 @@ uv run flashdreams-app t2v-app webrtc --prompt "A waterfall" A compatible package must expose an importable module with: -- `create_app(config: flashdreams_app.AppConfig)`, returning a +- `create_app_spec(config: flashdreams_app.AppConfig)`, returning a `PipelineAppSpec` with a `StreamInferencePipelineConfig`, initial conditioning, presentation metadata, step count, and a `PipelineContract` cache initializer. - Optionally, `add_arguments(parser)` adds provider-specific flags. -The host owns process/distributed initialization, pipeline setup, execution -options, runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, -and WebRTC serving. Providers only select/configure a pipeline and describe how -global conditioning initializes its cache. +The host owns process/distributed initialization, pipeline setup, +runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, and +WebRTC serving. Providers only select/configure a pipeline and describe how +global conditioning initializes its cache. Pipeline-specific execution options, +including compilation and CUDA graphs, remain encapsulated by the pipeline +config and its `setup()` implementation. diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py index 2d3e45aa9..a8ed56da7 100644 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -36,12 +36,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("mode", choices=("mp4", "webrtc")) parser.add_argument("--output", type=Path, help="MP4 path (required for mp4)") parser.add_argument("--device", default="cuda", help="Runtime device") - parser.add_argument( - "--compile", action=argparse.BooleanOptionalAction, default=None - ) - parser.add_argument( - "--cuda-graph", action=argparse.BooleanOptionalAction, default=None - ) 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") parser.add_argument("--warmup-chunks", type=int, default=0) @@ -113,20 +107,20 @@ def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: options["world_rank"] = environment.world_rank options["world_size"] = environment.world_size - factory = getattr(provider, "create_app", None) + factory = getattr(provider, "create_app_spec", None) if not callable(factory): - raise TypeError(f"Provider {args.provider!r} must define create_app(config).") + raise TypeError( + f"Provider {args.provider!r} must define create_app_spec(config)." + ) spec = factory(AppConfig(options=options)) if not isinstance(spec, PipelineAppSpec): raise TypeError( - f"Provider {args.provider!r} create_app() returned " + f"Provider {args.provider!r} create_app_spec() returned " f"{type(spec).__name__}, expected PipelineAppSpec." ) runtime = PipelineAppRuntime( spec=spec, device=environment.device, - compile=args.compile, - cuda_graph=args.cuda_graph, ) if args.mode == "webrtc": return _run_webrtc(runtime=runtime, args=args, environment=environment) diff --git a/apps/flashdreams_app/flashdreams_app/runtime.py b/apps/flashdreams_app/flashdreams_app/runtime.py index f5d9546a3..20f05dff4 100644 --- a/apps/flashdreams_app/flashdreams_app/runtime.py +++ b/apps/flashdreams_app/flashdreams_app/runtime.py @@ -10,8 +10,6 @@ import torch -from flashdreams.infra.config import derive_config -from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.runtime import ( InferenceInput, InferenceSession, @@ -31,15 +29,8 @@ def __init__( *, spec: PipelineAppSpec, device: str, - compile: bool | None = None, - cuda_graph: bool | None = None, ) -> None: - pipeline_config = _with_execution_options( - spec.pipeline_config, - compile=compile, - cuda_graph=cuda_graph, - ) - self.pipeline = pipeline_config.setup().to(device).eval() + self.pipeline = spec.pipeline_config.setup().to(device).eval() self.metadata = spec.metadata self.initial_input = spec.initial_input self._spec = spec @@ -158,25 +149,6 @@ def close(self) -> None: close() -def _with_execution_options( - pipeline: StreamInferencePipelineConfig, - *, - compile: bool | None, - cuda_graph: bool | None, -) -> StreamInferencePipelineConfig: - transformer: dict[str, object] = {} - if compile is not None: - transformer["compile_network"] = compile - if cuda_graph is not None: - transformer["use_cuda_graph"] = cuda_graph - if not transformer: - return pipeline - return derive_config( - pipeline, - diffusion_model={"transformer": transformer}, - ) - - def _metrics(value: object) -> Mapping[str, float | int]: if value is None: return {} diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py index 261fbac22..e7aafe555 100644 --- a/apps/flashdreams_app/tests/test_cli.py +++ b/apps/flashdreams_app/tests/test_cli.py @@ -4,13 +4,12 @@ from __future__ import annotations import argparse -from types import ModuleType, SimpleNamespace +from types import ModuleType from typing import Any, cast import pytest import torch from flashdreams_app import ( - PipelineAppRuntime, PipelineAppSpec, PipelineContract, RuntimeMetadata, @@ -34,7 +33,7 @@ def close(self) -> None: class Pipeline: def __init__(self, config: object) -> None: - del config + assert config is pipeline_config calls.append("pipeline.init") def to(self, device: str) -> "Pipeline": @@ -78,7 +77,7 @@ def initialize_cache(pipeline: object, inputs: InferenceInput) -> object: provider = ModuleType("fake_app") setattr( provider, - "create_app", + "create_app_spec", lambda config: PipelineAppSpec( pipeline_config=pipeline_config, contract=PipelineContract(initialize_cache=initialize_cache), @@ -134,6 +133,12 @@ def test_host_exposes_only_supported_output_modes() -> None: assert mode_action.choices == ("mp4", "webrtc") +def test_host_does_not_expose_pipeline_execution_options() -> None: + destinations = {action.dest for action in cli.build_parser()._actions} + assert "compile" not in destinations + assert "cuda_graph" not in destinations + + def test_webrtc_path_owns_serving_options_and_runtime_close( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -189,60 +194,3 @@ def serve(**kwargs: object) -> None: assert options.host == "127.0.0.1" assert options.port == 9000 assert options.encoder_gop == 24 - - -def test_host_owns_execution_options() -> None: - destinations = {action.dest for action in cli.build_parser()._actions} - assert {"compile", "cuda_graph"} <= destinations - - -def test_host_applies_execution_options_without_mutating_provider_spec() -> None: - configured_pipeline: StreamInferencePipelineConfig | None = None - - class Pipeline: - def __init__(self, config: StreamInferencePipelineConfig) -> None: - nonlocal configured_pipeline - configured_pipeline = config - - def to(self, device: str) -> "Pipeline": - assert device == "cpu" - return self - - def eval(self) -> "Pipeline": - return self - - transformer = SimpleNamespace(compile_network=False, use_cuda_graph=False) - pipeline_config = StreamInferencePipelineConfig( - _target=cast(Any, Pipeline), - name="fake", - diffusion_model=cast(Any, SimpleNamespace(transformer=transformer)), - ) - spec = PipelineAppSpec( - pipeline_config=pipeline_config, - contract=PipelineContract(initialize_cache=lambda pipeline, inputs: object()), - metadata=RuntimeMetadata( - model_id="fake", - fps=24, - output_layout="tchw", - video_width=64, - video_height=64, - ), - initial_input=InferenceInput(), - total_steps=1, - ) - - runtime = PipelineAppRuntime( - spec=spec, - device="cpu", - compile=True, - cuda_graph=True, - ) - try: - assert configured_pipeline is not None - configured_transformer = configured_pipeline.diffusion_model.transformer - assert getattr(configured_transformer, "compile_network") is True - assert getattr(configured_transformer, "use_cuda_graph") is True - assert transformer.compile_network is False - assert transformer.use_cuda_graph is False - finally: - runtime.close() diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py index d4aff7bd6..490a3d4d1 100644 --- a/apps/t2v_app/t2v_app/__init__.py +++ b/apps/t2v_app/t2v_app/__init__.py @@ -5,6 +5,6 @@ from flashdreams.core.pipeline_presets import PipelineProvider -from .provider import add_arguments, create_app +from .provider import add_arguments, create_app_spec -__all__ = ["PipelineProvider", "add_arguments", "create_app"] +__all__ = ["PipelineProvider", "add_arguments", "create_app_spec"] diff --git a/apps/t2v_app/t2v_app/provider.py b/apps/t2v_app/t2v_app/provider.py index 4fb0f3f3e..7fe0eaf01 100644 --- a/apps/t2v_app/t2v_app/provider.py +++ b/apps/t2v_app/t2v_app/provider.py @@ -49,7 +49,7 @@ def add_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument("--fps", type=int) -def create_app(config: AppConfig) -> PipelineAppSpec: +def create_app_spec(config: AppConfig) -> PipelineAppSpec: """Describe a T2V pipeline application without constructing its runtime.""" options = config.options catalog = load_preset_catalog(_optional_path(options.get("preset_config"))) @@ -200,4 +200,4 @@ def _optional_string(value: object) -> str | None: return value.strip() -__all__ = ["add_arguments", "create_app"] +__all__ = ["add_arguments", "create_app_spec"] diff --git a/apps/t2v_app/tests/test_provider.py b/apps/t2v_app/tests/test_provider.py index 73c42dc3c..1f3eaaca2 100644 --- a/apps/t2v_app/tests/test_provider.py +++ b/apps/t2v_app/tests/test_provider.py @@ -33,7 +33,7 @@ def test_t2v_provider_registers_model_options() -> None: assert not hasattr(args, "compile") -def test_create_app_returns_data_without_constructing_pipeline( +def test_create_app_spec_returns_data_without_constructing_pipeline( monkeypatch: pytest.MonkeyPatch, ) -> None: pipeline_constructed = False @@ -78,7 +78,7 @@ def create_pipeline_config( monkeypatch.setattr(provider, "load_preset_catalog", lambda _: catalog) monkeypatch.setattr(provider, "load_pipeline_provider", lambda _: Provider()) - created = provider.create_app( + created = provider.create_app_spec( AppConfig( options={ "preset_config": None, From b5d9b27e2ec158a24dedd96f5e6d11bcc84bc672 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 02:56:35 +0000 Subject: [PATCH 04/10] Define FlashDreams app provider interface Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 5 ++- apps/flashdreams_app/README.md | 6 ++- .../flashdreams_app/__init__.py | 2 + apps/flashdreams_app/flashdreams_app/cli.py | 32 +++++++++------- .../flashdreams_app/contracts.py | 14 +++++++ apps/flashdreams_app/tests/test_cli.py | 37 ++++++++++++++++++- apps/t2v_app/tests/test_provider.py | 7 +++- 7 files changed, 85 insertions(+), 18 deletions(-) diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md index 6808b6647..9e887f275 100644 --- a/PR_DESCRIPTION.md +++ b/PR_DESCRIPTION.md @@ -53,11 +53,14 @@ When `--preset-id` is omitted, `t2v-app` uses the catalog's - Add the `flashdreams-app` workspace package and console entrypoint. - Define a minimal provider boundary: `create_app_spec(AppConfig) -> PipelineAppSpec`. +- Require provider modules to conform to `AppProvider` with `add_arguments()` + and `create_app_spec()`. - Add the host-owned `PipelineAppRuntime` and `PipelineAppSession`. - Keep pipeline setup, `generate`/`finalize`, step tracking, cache release, and runtime closure in the host. - Add host-owned MP4 and WebRTC presentation paths without a runtime adapter. -- Keep provider-specific CLI arguments optional through `add_arguments(parser)`. +- Let providers register custom CLI arguments through `add_arguments(parser)`; + providers without custom arguments use a no-op implementation. - Keep pipeline-specific execution behavior encapsulated by the selected pipeline config and its `setup()` implementation. diff --git a/apps/flashdreams_app/README.md b/apps/flashdreams_app/README.md index c8b697291..3e2d78177 100644 --- a/apps/flashdreams_app/README.md +++ b/apps/flashdreams_app/README.md @@ -14,11 +14,15 @@ uv run flashdreams-app t2v-app webrtc --prompt "A waterfall" A compatible package must expose an importable module with: +- `add_arguments(parser)`, which registers provider-specific flags. Providers + without custom flags implement this as a no-op. - `create_app_spec(config: flashdreams_app.AppConfig)`, returning a `PipelineAppSpec` with a `StreamInferencePipelineConfig`, initial conditioning, presentation metadata, step count, and a `PipelineContract` cache initializer. -- Optionally, `add_arguments(parser)` adds provider-specific flags. + +The module structurally conforms to `flashdreams_app.AppProvider`; the host +validates this contract when it loads the installed provider. The host owns process/distributed initialization, pipeline setup, runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, and diff --git a/apps/flashdreams_app/flashdreams_app/__init__.py b/apps/flashdreams_app/flashdreams_app/__init__.py index 9a17f6ab0..c5d28caf9 100644 --- a/apps/flashdreams_app/flashdreams_app/__init__.py +++ b/apps/flashdreams_app/flashdreams_app/__init__.py @@ -5,6 +5,7 @@ from .contracts import ( AppConfig, + AppProvider, AppRuntime, PipelineAppSpec, PipelineContract, @@ -15,6 +16,7 @@ __all__ = [ "AppConfig", + "AppProvider", "AppRuntime", "PipelineAppRuntime", "PipelineAppSpec", diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py index a8ed56da7..c88b793ec 100644 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -10,7 +10,6 @@ from contextlib import ExitStack from importlib import metadata from pathlib import Path -from types import ModuleType from typing import Sequence import torch @@ -21,7 +20,7 @@ ) from flashdreams.runtime.output import OutputArtifact -from .contracts import AppConfig, AppRuntime, PipelineAppSpec +from .contracts import AppConfig, AppProvider, AppRuntime, PipelineAppSpec from .outputs import FileOutput from .runtime import PipelineAppRuntime from .webrtc import WebRTCOptions, serve_webrtc @@ -49,8 +48,8 @@ def build_parser() -> argparse.ArgumentParser: return parser -def load_provider(distribution_name: str) -> ModuleType: - """Load an installed provider after verifying its distribution is present.""" +def load_provider(distribution_name: str) -> AppProvider: + """Load an installed provider that satisfies the host contract.""" try: distribution = metadata.distribution(distribution_name) except metadata.PackageNotFoundError as exc: @@ -65,12 +64,24 @@ def load_provider(distribution_name: str) -> ModuleType: if distribution.metadata["Name"] in distributions ] candidates.append(distribution_name.replace("-", "_")) + incompatible: list[str] = [] for candidate in dict.fromkeys(candidates): try: - return importlib.import_module(candidate) + module = importlib.import_module(candidate) except ModuleNotFoundError as exc: if exc.name != candidate: raise + continue + if isinstance(module, AppProvider): + return module + incompatible.append(candidate) + if incompatible: + names = ", ".join(repr(name) for name in incompatible) + raise TypeError( + f"Provider distribution {distribution_name!r} exposes module(s) " + f"{names}, but none satisfy AppProvider. Providers must define " + "add_arguments(parser) and create_app_spec(config)." + ) raise ValueError( f"Provider {distribution_name!r} does not expose an importable Python module." ) @@ -94,9 +105,7 @@ def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: parser.parse_args(argv) return () provider = load_provider(provider_args.provider) - add_arguments = getattr(provider, "add_arguments", None) - if callable(add_arguments): - add_arguments(parser) + provider.add_arguments(parser) args = parser.parse_args(argv) if args.mode == "mp4" and args.output is None: parser.error("--output is required for mp4 mode") @@ -107,12 +116,7 @@ def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: options["world_rank"] = environment.world_rank options["world_size"] = environment.world_size - factory = getattr(provider, "create_app_spec", None) - if not callable(factory): - raise TypeError( - f"Provider {args.provider!r} must define create_app_spec(config)." - ) - spec = factory(AppConfig(options=options)) + spec = provider.create_app_spec(AppConfig(options=options)) if not isinstance(spec, PipelineAppSpec): raise TypeError( f"Provider {args.provider!r} create_app_spec() returned " diff --git a/apps/flashdreams_app/flashdreams_app/contracts.py b/apps/flashdreams_app/flashdreams_app/contracts.py index f7f1be2a3..59b8bbc34 100644 --- a/apps/flashdreams_app/flashdreams_app/contracts.py +++ b/apps/flashdreams_app/flashdreams_app/contracts.py @@ -5,6 +5,7 @@ from __future__ import annotations +import argparse from collections.abc import Callable, Mapping from dataclasses import dataclass, field from types import MappingProxyType @@ -96,6 +97,19 @@ def __post_init__(self) -> None: ) +@runtime_checkable +class AppProvider(Protocol): + """Provider module contract consumed by the application host.""" + + def add_arguments(self, parser: argparse.ArgumentParser) -> None: + """Register provider-specific arguments on the host parser.""" + ... + + def create_app_spec(self, config: AppConfig) -> PipelineAppSpec: + """Create a declarative application specification.""" + ... + + def require_pipeline_config( config: object, *, expected_name: str | None = None ) -> StreamInferencePipelineConfig: diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py index e7aafe555..774f125af 100644 --- a/apps/flashdreams_app/tests/test_cli.py +++ b/apps/flashdreams_app/tests/test_cli.py @@ -4,12 +4,13 @@ from __future__ import annotations import argparse -from types import ModuleType +from types import ModuleType, SimpleNamespace from typing import Any, cast import pytest import torch from flashdreams_app import ( + AppProvider, PipelineAppSpec, PipelineContract, RuntimeMetadata, @@ -75,6 +76,12 @@ def initialize_cache(pipeline: object, inputs: InferenceInput) -> object: return Cache() provider = ModuleType("fake_app") + + def add_arguments(parser: argparse.ArgumentParser) -> None: + del parser + calls.append("provider.add_arguments") + + setattr(provider, "add_arguments", add_arguments) setattr( provider, "create_app_spec", @@ -111,6 +118,7 @@ def close(self) -> tuple[object, ...]: monkeypatch.setattr(cli, "FileOutput", Output) cli.run(["fake-app", "mp4", "--device", "cpu", "--output", "result.mp4"]) assert calls == [ + "provider.add_arguments", "pipeline.init", "pipeline.to", "pipeline.eval", @@ -126,6 +134,33 @@ def close(self) -> tuple[object, ...]: ] +def test_app_provider_protocol_requires_both_methods() -> None: + provider = ModuleType("provider") + setattr(provider, "add_arguments", lambda parser: None) + setattr(provider, "create_app_spec", lambda config: None) + assert isinstance(provider, AppProvider) + + delattr(provider, "add_arguments") + assert not isinstance(provider, AppProvider) + + +def test_load_provider_rejects_module_outside_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = ModuleType("invalid_provider") + distribution = SimpleNamespace(metadata={"Name": "invalid-app"}) + monkeypatch.setattr(cli.metadata, "distribution", lambda name: distribution) + monkeypatch.setattr( + cli.metadata, + "packages_distributions", + lambda: {"invalid_provider": ["invalid-app"]}, + ) + monkeypatch.setattr(cli.importlib, "import_module", lambda name: module) + + with pytest.raises(TypeError, match="none satisfy AppProvider"): + cli.load_provider("invalid-app") + + def test_host_exposes_only_supported_output_modes() -> None: mode_action = next( action for action in cli.build_parser()._actions if action.dest == "mode" diff --git a/apps/t2v_app/tests/test_provider.py b/apps/t2v_app/tests/test_provider.py index 1f3eaaca2..78c3d307e 100644 --- a/apps/t2v_app/tests/test_provider.py +++ b/apps/t2v_app/tests/test_provider.py @@ -9,7 +9,8 @@ from typing import Any, cast import pytest -from flashdreams_app import AppConfig, PipelineAppSpec +import t2v_app +from flashdreams_app import AppConfig, AppProvider, PipelineAppSpec from t2v_app import provider from t2v_app.presets import ( PipelinePreset, @@ -22,6 +23,10 @@ pytestmark = pytest.mark.ci_cpu +def test_provider_module_conforms_to_host_contract() -> None: + assert isinstance(t2v_app, AppProvider) + + def test_t2v_provider_registers_model_options() -> None: parser = argparse.ArgumentParser() provider.add_arguments(parser) From c03fcc032718742285dce4b0d3d9353ba52253d6 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 03:03:06 +0000 Subject: [PATCH 05/10] Simplify T2V app provider surface Signed-off-by: Gangzheng Tong --- apps/t2v_app/t2v_app/__init__.py | 4 +--- apps/t2v_app/t2v_app/provider.py | 39 ++++++++++++++++++++++++++------ 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py index 490a3d4d1..9a8680c96 100644 --- a/apps/t2v_app/t2v_app/__init__.py +++ b/apps/t2v_app/t2v_app/__init__.py @@ -3,8 +3,6 @@ """Text-to-video pipeline application for ``flashdreams-app``.""" -from flashdreams.core.pipeline_presets import PipelineProvider - from .provider import add_arguments, create_app_spec -__all__ = ["PipelineProvider", "add_arguments", "create_app_spec"] +__all__ = ["add_arguments", "create_app_spec"] diff --git a/apps/t2v_app/t2v_app/provider.py b/apps/t2v_app/t2v_app/provider.py index 7fe0eaf01..753004cd1 100644 --- a/apps/t2v_app/t2v_app/provider.py +++ b/apps/t2v_app/t2v_app/provider.py @@ -20,9 +20,10 @@ from flashdreams.core.pipeline_presets import load_pipeline_provider from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.runtime import InferenceInput -from .presets import RuntimePresetOptions, load_preset_catalog +from .presets import PipelinePreset, RuntimePresetOptions, load_preset_catalog FIELD_PROMPT = "prompt" FIELD_TOTAL_BLOCKS = "total_blocks" @@ -52,24 +53,48 @@ def add_arguments(parser: argparse.ArgumentParser) -> None: def create_app_spec(config: AppConfig) -> PipelineAppSpec: """Describe a T2V pipeline application without constructing its runtime.""" options = config.options + preset_id, preset = _resolve_preset(options) + scenario = _scenario(options, preset.runtime) + pipeline_config = _create_pipeline_config(preset_id, preset) + return _build_app_spec(pipeline_config, scenario, preset.runtime) + + +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"))) - preset_id, preset = catalog.resolve(_optional_string(options.get("preset_id"))) - provider = load_pipeline_provider(preset.provider) - pipeline_config = require_pipeline_config( - provider.create_pipeline_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) + return require_pipeline_config( + pipeline_provider.create_pipeline_config( preset_id=preset_id, options=preset.pipeline, ), expected_name=preset_id, ) - scenario = _scenario(options, preset.runtime) + + +def _build_app_spec( + pipeline_config: StreamInferencePipelineConfig, + scenario: Mapping[str, object], + runtime_options: RuntimePresetOptions, +) -> PipelineAppSpec: + """Build the host-facing application spec from resolved T2V data.""" return PipelineAppSpec( pipeline_config=pipeline_config, contract=PipelineContract(initialize_cache=_initialize_cache), metadata=RuntimeMetadata( model_id="t2v-app", fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), - output_layout=preset.runtime.output_layout, + output_layout=runtime_options.output_layout, video_width=_required_int( scenario[FIELD_PIXEL_WIDTH], name=FIELD_PIXEL_WIDTH ), From 9c4894ec424eed84a61d92eb70a57f1efaa9118a Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 03:09:59 +0000 Subject: [PATCH 06/10] Simplify WebRTC prototype options Signed-off-by: Gangzheng Tong --- apps/flashdreams_app/flashdreams_app/cli.py | 18 ++------ .../flashdreams_app/flashdreams_app/webrtc.py | 46 ++++--------------- apps/flashdreams_app/tests/test_cli.py | 21 ++++++--- apps/flashdreams_app/tests/test_webrtc.py | 13 +++--- 4 files changed, 33 insertions(+), 65 deletions(-) diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py index c88b793ec..83221f986 100644 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -37,14 +37,6 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--device", default="cuda", help="Runtime device") 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") - parser.add_argument("--warmup-chunks", type=int, default=0) - parser.add_argument("--warmup-timeout-s", type=float, default=600.0) - parser.add_argument("--client-liveness-timeout-s", type=float, default=30.0) - parser.add_argument( - "--encoder-backend", choices=("auto", "default", "nvenc"), default="auto" - ) - parser.add_argument("--encoder-bitrate-bps", type=int, default=6_000_000) - parser.add_argument("--encoder-gop", type=int) return parser @@ -159,14 +151,8 @@ def _run_webrtc( options=WebRTCOptions( host=args.host, port=args.port, - warmup_chunks=args.warmup_chunks, - warmup_timeout_s=args.warmup_timeout_s, - client_liveness_timeout_s=args.client_liveness_timeout_s, - device=environment.device, - encoder_backend=args.encoder_backend, - encoder_bitrate_bps=args.encoder_bitrate_bps, - encoder_gop=args.encoder_gop or int(runtime.metadata.fps), ), + device=environment.device, world_rank=environment.world_rank, ) finally: @@ -239,5 +225,7 @@ def _initialize_environment(device: str) -> _Environment: def main() -> None: """Console-script entry point.""" artifacts = run() + # File modes return persistent artifacts whose URI is the output location; + # live modes such as WebRTC return no artifacts. for artifact in artifacts: print(artifact.uri) diff --git a/apps/flashdreams_app/flashdreams_app/webrtc.py b/apps/flashdreams_app/flashdreams_app/webrtc.py index 32f6d8a4f..fb897c956 100644 --- a/apps/flashdreams_app/flashdreams_app/webrtc.py +++ b/apps/flashdreams_app/flashdreams_app/webrtc.py @@ -26,31 +26,17 @@ from .contracts import AppRuntime +# TODO: Move this contract into shared FlashDreams and expand it when the +# serving API needs caller-configurable transport and encoder tuning. @dataclass(frozen=True, slots=True) class WebRTCOptions: - """Presentation settings owned by ``flashdreams-app``.""" + """Minimal WebRTC bind settings for the application prototype.""" host: str - port: int - warmup_chunks: int - warmup_timeout_s: float - client_liveness_timeout_s: float - device: str - encoder_backend: str - encoder_bitrate_bps: int - encoder_gop: int - + """Server bind address.""" -@dataclass(frozen=True, slots=True) -class _WebRTCRuntimeConfig: - video_width: int - video_height: int - warmup_chunks: int - warmup_timeout_s: float - device: str - encoder_backend: str - encoder_bitrate_bps: int - encoder_gop: int + port: int + """Server bind port.""" class _InputProvider: @@ -79,7 +65,7 @@ def close(self) -> None: def serve_webrtc( - *, runtime: AppRuntime, options: WebRTCOptions, world_rank: int + *, runtime: AppRuntime, options: WebRTCOptions, device: str, world_rank: int ) -> object: """Serve an application runtime through the shared WebRTC stack.""" metadata = runtime.metadata @@ -89,16 +75,13 @@ def serve_webrtc( fps=int(metadata.fps), video_width=metadata.video_width, video_height=metadata.video_height, - warmup_chunks=options.warmup_chunks, - warmup_timeout_s=options.warmup_timeout_s, - client_liveness_timeout_s=options.client_liveness_timeout_s, preload_name=metadata.model_id, ) spec = DemoSpec( model_id=metadata.model_id, input_mode="webrtc", output=output, - config=InferenceConfig(model_id=metadata.model_id, device=options.device), + config=InferenceConfig(model_id=metadata.model_id, device=device), ) scenario = PreparedScenario(initial_inputs=runtime.initial_input) @@ -111,23 +94,14 @@ def create_model_input_provider( manager = BaseWebRTCSessionManager( runtime=runtime, - runtime_config=_WebRTCRuntimeConfig( - video_width=metadata.video_width, - video_height=metadata.video_height, - warmup_chunks=options.warmup_chunks, - warmup_timeout_s=options.warmup_timeout_s, - device=options.device, - encoder_backend=options.encoder_backend, - encoder_bitrate_bps=options.encoder_bitrate_bps, - encoder_gop=options.encoder_gop, - ), + runtime_config=output, fps=int(metadata.fps), identity=metadata.model_id, shared_host=RuntimeHost(runtime), shared_spec=spec, shared_scenario=scenario, shared_model_input_provider_factory=create_model_input_provider, - client_liveness_timeout_s=options.client_liveness_timeout_s, + client_liveness_timeout_s=output.client_liveness_timeout_s, ) return serve_webrtc_demo( output=output, diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py index 774f125af..0c5aee0f1 100644 --- a/apps/flashdreams_app/tests/test_cli.py +++ b/apps/flashdreams_app/tests/test_cli.py @@ -174,6 +174,19 @@ def test_host_does_not_expose_pipeline_execution_options() -> None: assert "cuda_graph" not in destinations +def test_host_exposes_only_minimal_webrtc_options() -> None: + destinations = {action.dest for action in cli.build_parser()._actions} + assert {"host", "port"} <= destinations + assert { + "warmup_chunks", + "warmup_timeout_s", + "client_liveness_timeout_s", + "encoder_backend", + "encoder_bitrate_bps", + "encoder_gop", + }.isdisjoint(destinations) + + def test_webrtc_path_owns_serving_options_and_runtime_close( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -211,12 +224,6 @@ def serve(**kwargs: object) -> None: args=argparse.Namespace( host="127.0.0.1", port=9000, - warmup_chunks=2, - warmup_timeout_s=30.0, - client_liveness_timeout_s=10.0, - encoder_backend="default", - encoder_bitrate_bps=1_000_000, - encoder_gop=None, ), environment=cli._Environment(device="cpu", world_rank=0, world_size=1), ) @@ -228,4 +235,4 @@ def serve(**kwargs: object) -> None: assert isinstance(options, cli.WebRTCOptions) assert options.host == "127.0.0.1" assert options.port == 9000 - assert options.encoder_gop == 24 + assert captured["device"] == "cpu" diff --git a/apps/flashdreams_app/tests/test_webrtc.py b/apps/flashdreams_app/tests/test_webrtc.py index 008e68d60..6be395dbe 100644 --- a/apps/flashdreams_app/tests/test_webrtc.py +++ b/apps/flashdreams_app/tests/test_webrtc.py @@ -63,25 +63,24 @@ def fake_serve(**kwargs: object) -> str: options=webrtc.WebRTCOptions( host="127.0.0.1", port=8080, - warmup_chunks=0, - warmup_timeout_s=30.0, - client_liveness_timeout_s=30.0, - device="cpu", - encoder_backend="default", - encoder_bitrate_bps=1_000_000, - encoder_gop=16, ), + 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.warmup_chunks == 0 + assert output.client_liveness_timeout_s == 30.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.runtime_config is output assert session_manager._shared_scenario is not None assert session_manager._shared_scenario.initial_inputs is runtime.initial_input assert callable(session_manager._shared_model_input_provider_factory) From ec35d93bc552b405a941b226d2c10318d17b543a Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 04:38:48 +0000 Subject: [PATCH 07/10] Simplify FlashDreams application runtime contract Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 50 +++- apps/flashdreams_app/README.md | 106 +++++++- .../flashdreams_app/__init__.py | 20 +- apps/flashdreams_app/flashdreams_app/cli.py | 194 +++++++++++---- .../flashdreams_app/contracts.py | 196 ++++++++------- .../flashdreams_app/runtime.py | 39 ++- .../flashdreams_app/flashdreams_app/webrtc.py | 62 +++-- apps/flashdreams_app/tests/test_cli.py | 227 ++++++++++++++---- apps/flashdreams_app/tests/test_webrtc.py | 26 +- apps/t2v_app/README.md | 24 +- apps/t2v_app/pyproject.toml | 5 + apps/t2v_app/t2v_app/__init__.py | 4 +- apps/t2v_app/t2v_app/presets.py | 37 ++- apps/t2v_app/t2v_app/provider.py | 141 +++++++---- apps/t2v_app/tests/test_pipeline_provider.py | 37 ++- apps/t2v_app/tests/test_provider.py | 69 ++++-- 16 files changed, 871 insertions(+), 366 deletions(-) diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md index 9e887f275..f39886847 100644 --- a/PR_DESCRIPTION.md +++ b/PR_DESCRIPTION.md @@ -7,7 +7,36 @@ services. The host owns process initialization, pipeline construction, the runtime/session lifecycle, autoregressive stepping, finalization, cleanup, and presentation. The T2V package only selects a pipeline preset and supplies the -model-specific conditioning and cache-initialization contract. +model-specific conditioning and cache-initialization callback. + +## High-level design + +```text +uv run flashdreams-app t2v-app {mp4 | webrtc} + | + v ++----------------------+ request spec +----------------------+ +| flashdreams_app | --------------------> | t2v_app | +| generic host | <-------------------- | declarative provider | ++----------+-----------+ AppSpec +----------+-----------+ + | | + | constructs and drives | reads + v v ++----------------------+ +------------------------+ +| FlashDreams runtime | | pipeline_presets.yaml | +| API (black box) | | + PipelineProvider | ++----------+-----------+ +------------------------+ + | + | video chunks + v + +----+----+ + | | + v v ++----------+ +----------+ +| MP4 file | | WebRTC | +| artifact | | stream | ++----------+ +----------+ +``` ## Entrypoint examples @@ -52,23 +81,28 @@ When `--preset-id` is omitted, `t2v-app` uses the catalog's - Add the `flashdreams-app` workspace package and console entrypoint. - Define a minimal provider boundary: - `create_app_spec(AppConfig) -> PipelineAppSpec`. -- Require provider modules to conform to `AppProvider` with `add_arguments()` - and `create_app_spec()`. + `create_app_spec(AppRequest) -> AppSpec`. +- Keep the provider surface data-first: a mode-independent pipeline spec plus + an `AppConfig` for presentation and mode-specific MP4 or WebRTC run data. +- Require provider modules to conform to `AppProvider` with + `parse_options(parser, argv)` and `create_app_spec(request)`. - Add the host-owned `PipelineAppRuntime` and `PipelineAppSession`. - Keep pipeline setup, `generate`/`finalize`, step tracking, cache release, and runtime closure in the host. - Add host-owned MP4 and WebRTC presentation paths without a runtime adapter. -- Let providers register custom CLI arguments through `add_arguments(parser)`; - providers without custom arguments use a no-op implementation. +- Type both presentation paths directly against the shared `InferenceRuntime` + contract. +- Parse only the provider and presentation mode in the host, then let the + provider extend the mode-specific parser and parse all remaining arguments + through `parse_options(parser, argv)`. - Keep pipeline-specific execution behavior encapsulated by the selected pipeline config and its `setup()` implementation. ## T2V provider and presets - Add the `t2v-app` workspace package. -- Describe T2V through a `PipelineAppSpec` rather than implementing another - runtime or session. +- Describe T2V through an `AppSpec` rather than implementing another runtime + or session. - Add packaged YAML presets for causal-forcing and self-forcing WAN pipelines. - Resolve pipeline providers directly from the YAML catalog without depending on the runner registry. diff --git a/apps/flashdreams_app/README.md b/apps/flashdreams_app/README.md index 3e2d78177..27ab31da7 100644 --- a/apps/flashdreams_app/README.md +++ b/apps/flashdreams_app/README.md @@ -12,21 +12,103 @@ uv run flashdreams-app t2v-app webrtc --prompt "A waterfall" ## Provider contract -A compatible package must expose an importable module with: +A compatible package must expose an importable module with exactly two public +entry points: -- `add_arguments(parser)`, which registers provider-specific flags. Providers - without custom flags implement this as a no-op. -- `create_app_spec(config: flashdreams_app.AppConfig)`, returning a - `PipelineAppSpec` with a `StreamInferencePipelineConfig`, initial - conditioning, presentation metadata, step count, and a `PipelineContract` - cache initializer. +- `parse_options(parser, argv)`, which extends the selected mode's host parser + with provider-specific flags, parses the remaining arguments, and returns the + parsed values as a mapping. +- `create_app_spec(request: flashdreams_app.AppRequest)`, returning a + mode-aware `AppSpec` containing a shared `PipelineAppSpec` and either an + `Mp4RunSpec` or `WebRTCRunSpec`. The module structurally conforms to `flashdreams_app.AppProvider`; the host validates this contract when it loads the installed provider. -The host owns process/distributed initialization, pipeline setup, +`PipelineAppSpec` contains only mode-independent runtime data: + +| Provider-supplied field | Purpose | +|---|---| +| `pipeline_config` | A `StreamInferencePipelineConfig`; the host calls `setup()` and owns the resulting pipeline. | +| `initialize_cache` | A callback receiving the constructed pipeline and session input; it returns the session cache. | + +`AppSpec.config` is an `AppConfig` containing the model identity, video +dimensions, frame rate, and output tensor layout used by host-owned +presentation. `AppRequest` contains only the selected mode and parsed invocation +options passed into the provider. + +Invocation data belongs to the selected presentation mode: + +| Run spec | Required data | +|---|---| +| `Mp4RunSpec` | `initial_input` and the finite `total_steps` written to the file. | +| `WebRTCRunSpec` | `initial_input` for each live session; no finite step count. | + +The following is a complete provider skeleton: + +```python +import argparse +from collections.abc import Mapping, Sequence +from typing import Any + +from flashdreams_app import ( + AppConfig, + AppRequest, + AppSpec, + Mp4RunSpec, + PipelineAppSpec, + WebRTCRunSpec, +) +from flashdreams.runtime import InferenceInput +from my_model.config import MY_PIPELINE_CONFIG + + +def parse_options( + parser: argparse.ArgumentParser, + argv: Sequence[str], +) -> Mapping[str, Any]: + parser.add_argument("--prompt", required=True) + return vars(parser.parse_args(argv)) + + +def create_app_spec(request: AppRequest) -> AppSpec: + prompt = str(request.options["prompt"]) + initial_input = InferenceInput(global_conditioning={"prompt": prompt}) + if request.mode == "mp4": + run = Mp4RunSpec(initial_input=initial_input, total_steps=4) + else: + run = WebRTCRunSpec(initial_input=initial_input) + return AppSpec( + config=AppConfig( + model_id="my-model", + fps=24, + output_layout="tchw", + video_width=832, + video_height=480, + ), + pipeline=PipelineAppSpec( + pipeline_config=MY_PIPELINE_CONFIG, + initialize_cache=_initialize_cache, + ), + run=run, + ) + + +def _initialize_cache(pipeline: Any, inputs: InferenceInput) -> object: + prompt = str(inputs.global_conditioning["prompt"]) + return pipeline.initialize_cache( + text=[prompt], + image=None, + height=60, + width=104, + ) +``` + +The host first parses only the provider and presentation mode. It then gives the +provider a parser containing that mode's host-owned options and all remaining +arguments. The host owns process/distributed initialization, pipeline setup, runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, and -WebRTC serving. Providers only select/configure a pipeline and describe how -global conditioning initializes its cache. Pipeline-specific execution options, -including compilation and CUDA graphs, remain encapsulated by the pipeline -config and its `setup()` implementation. +WebRTC serving. Providers only parse their options, select/configure a pipeline, +and describe how global conditioning initializes its cache. Pipeline-specific +execution options, including compilation and CUDA graphs, remain encapsulated +by the pipeline config and its `setup()` implementation. diff --git a/apps/flashdreams_app/flashdreams_app/__init__.py b/apps/flashdreams_app/flashdreams_app/__init__.py index c5d28caf9..fc79e7875 100644 --- a/apps/flashdreams_app/flashdreams_app/__init__.py +++ b/apps/flashdreams_app/flashdreams_app/__init__.py @@ -1,26 +1,24 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Host runtime and public provider contract for FlashDreams applications.""" +"""Public contract for FlashDreams application providers.""" from .contracts import ( AppConfig, AppProvider, - AppRuntime, + AppRequest, + AppSpec, + Mp4RunSpec, PipelineAppSpec, - PipelineContract, - RuntimeMetadata, - require_pipeline_config, + WebRTCRunSpec, ) -from .runtime import PipelineAppRuntime __all__ = [ "AppConfig", "AppProvider", - "AppRuntime", - "PipelineAppRuntime", + "AppRequest", + "AppSpec", + "Mp4RunSpec", "PipelineAppSpec", - "PipelineContract", - "RuntimeMetadata", - "require_pipeline_config", + "WebRTCRunSpec", ] diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py index 83221f986..59d42b9ee 100644 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ b/apps/flashdreams_app/flashdreams_app/cli.py @@ -8,35 +8,66 @@ import argparse import importlib from contextlib import ExitStack +from dataclasses import dataclass from importlib import metadata from pathlib import Path from typing import Sequence import torch +from flashdreams.runtime import InferenceInput, InferenceRuntime from flashdreams.runtime.demo.bootstrap import ( configure_logging, initialize_cuda_distributed, ) from flashdreams.runtime.output import OutputArtifact -from .contracts import AppConfig, AppProvider, AppRuntime, PipelineAppSpec +from .contracts import ( + AppConfig, + AppProvider, + AppRequest, + AppSpec, + Mp4RunSpec, + WebRTCRunSpec, +) from .outputs import FileOutput from .runtime import PipelineAppRuntime from .webrtc import WebRTCOptions, serve_webrtc -def build_parser() -> argparse.ArgumentParser: - """Build the host parser without importing a provider package.""" - parser = argparse.ArgumentParser(prog="flashdreams-app") - parser.add_argument( - "provider", help="Installed provider distribution, e.g. t2v-app" - ) - parser.add_argument("mode", choices=("mp4", "webrtc")) - parser.add_argument("--output", type=Path, help="MP4 path (required for mp4)") +@dataclass(frozen=True, slots=True) +class _ProviderAndMode: + """Top-level route parsed before loading an application provider.""" + + provider: str + """Installed provider distribution name.""" + + mode: str + """Selected presentation mode.""" + + remaining_argv: tuple[str, ...] + """Arguments delegated to the provider parser.""" + + +def build_parser(provider: str, mode: str) -> argparse.ArgumentParser: + """Build the selected mode's host-owned options parser. + + Args: + provider: Provider name displayed in command usage. + mode: Selected presentation mode. + + Returns: + Parser ready for the provider to extend and invoke. + """ + parser = argparse.ArgumentParser(prog=f"flashdreams-app {provider} {mode}") parser.add_argument("--device", default="cuda", help="Runtime device") - 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") + if mode == "mp4": + parser.add_argument("--output", type=Path, required=True, help="MP4 path") + elif 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") + else: + raise ValueError(f"Unsupported application mode: {mode!r}.") return parser @@ -72,7 +103,7 @@ def load_provider(distribution_name: str) -> AppProvider: raise TypeError( f"Provider distribution {distribution_name!r} exposes module(s) " f"{names}, but none satisfy AppProvider. Providers must define " - "add_arguments(parser) and create_app_spec(config)." + "parse_options(parser, argv) and create_app_spec(request)." ) raise ValueError( f"Provider {distribution_name!r} does not expose an importable Python module." @@ -88,50 +119,102 @@ def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: Returns: Artifacts produced by the selected path. WebRTC returns an empty tuple. """ - probe = argparse.ArgumentParser(add_help=False) - probe.add_argument("provider", nargs="?") - probe.add_argument("mode", nargs="?") - provider_args, _ = probe.parse_known_args(argv) - parser = build_parser() - if provider_args.provider is None: - parser.parse_args(argv) - return () - provider = load_provider(provider_args.provider) - provider.add_arguments(parser) - args = parser.parse_args(argv) - if args.mode == "mp4" and args.output is None: - parser.error("--output is required for mp4 mode") - + route = _parse_provider_and_mode(argv) + provider = load_provider(route.provider) + options = provider.parse_options( + build_parser(route.provider, route.mode), + route.remaining_argv, + ) + app_spec = _require_app_spec( + provider.create_app_spec(AppRequest(mode=route.mode, options=options)), + provider_name=route.provider, + mode=route.mode, + ) + args = argparse.Namespace(**options) environment = _initialize_environment(args.device) - options = vars(args).copy() - options["device"] = environment.device - options["world_rank"] = environment.world_rank - options["world_size"] = environment.world_size + runtime: InferenceRuntime = PipelineAppRuntime( + spec=app_spec.pipeline, + config=app_spec.config, + device=environment.device, + ) + return _launch_mode( + mode=route.mode, + runtime=runtime, + app_spec=app_spec, + args=args, + environment=environment, + ) + + +def _parse_provider_and_mode( + argv: Sequence[str] | None, +) -> _ProviderAndMode: + """Parse the provider and mode while preserving all remaining arguments.""" + parser = argparse.ArgumentParser(prog="flashdreams-app", add_help=False) + parser.add_argument("provider", help="Installed provider distribution") + parser.add_argument("mode", choices=("mp4", "webrtc")) + args, remaining_argv = parser.parse_known_args(argv) + return _ProviderAndMode( + provider=args.provider, + mode=args.mode, + remaining_argv=tuple(remaining_argv), + ) - spec = provider.create_app_spec(AppConfig(options=options)) - if not isinstance(spec, PipelineAppSpec): + +def _require_app_spec( + value: object, + *, + provider_name: str, + mode: str, +) -> AppSpec: + """Validate a provider result before constructing its runtime.""" + if not isinstance(value, AppSpec): raise TypeError( - f"Provider {args.provider!r} create_app_spec() returned " - f"{type(spec).__name__}, expected PipelineAppSpec." + f"Provider {provider_name!r} create_app_spec() returned " + f"{type(value).__name__}, expected AppSpec." ) - runtime = PipelineAppRuntime( - spec=spec, - device=environment.device, - ) - if args.mode == "webrtc": - return _run_webrtc(runtime=runtime, args=args, environment=environment) - if args.mode == "mp4": + if mode == "webrtc" and not isinstance(value.run, WebRTCRunSpec): + raise TypeError("WebRTC mode requires WebRTCRunSpec from the provider.") + if mode == "mp4" and not isinstance(value.run, Mp4RunSpec): + raise TypeError("MP4 mode requires Mp4RunSpec from the provider.") + return value + + +def _launch_mode( + *, + mode: str, + runtime: InferenceRuntime, + app_spec: AppSpec, + args: argparse.Namespace, + environment: "_Environment", +) -> tuple[OutputArtifact, ...]: + """Launch the host-owned presentation path selected by the user.""" + if mode == "webrtc": + assert isinstance(app_spec.run, WebRTCRunSpec) + return _run_webrtc( + runtime=runtime, + config=app_spec.config, + run_spec=app_spec.run, + args=args, + environment=environment, + ) + if mode == "mp4": + assert isinstance(app_spec.run, Mp4RunSpec) return _run_mp4( runtime=runtime, + config=app_spec.config, + run_spec=app_spec.run, output_path=args.output, environment=environment, ) - raise AssertionError(f"Unsupported presentation mode: {args.mode!r}.") + raise AssertionError(f"Unsupported presentation mode: {mode!r}.") def _run_webrtc( *, - runtime: AppRuntime, + runtime: InferenceRuntime, + config: AppConfig, + run_spec: WebRTCRunSpec, args: argparse.Namespace, environment: "_Environment", ) -> tuple[OutputArtifact, ...]: @@ -139,6 +222,8 @@ def _run_webrtc( Args: runtime: Initialized application runtime. + config: Model identity and video presentation configuration. + run_spec: Initial conditioning for live sessions. args: Parsed host and WebRTC arguments. environment: Initialized process and distributed environment. @@ -148,6 +233,8 @@ def _run_webrtc( try: serve_webrtc( runtime=runtime, + config=config, + initial_input=run_spec.initial_input, options=WebRTCOptions( host=args.host, port=args.port, @@ -162,7 +249,9 @@ def _run_webrtc( def _run_mp4( *, - runtime: AppRuntime, + runtime: InferenceRuntime, + config: AppConfig, + run_spec: Mp4RunSpec, output_path: Path, environment: "_Environment", ) -> tuple[OutputArtifact, ...]: @@ -170,6 +259,8 @@ def _run_mp4( Args: runtime: Initialized application runtime. + config: Model identity and video presentation configuration. + run_spec: Initial conditioning and finite generation length. output_path: Destination for the generated MP4. environment: Initialized process and distributed environment. @@ -181,8 +272,8 @@ def _run_mp4( resources.callback(runtime.close) output = FileOutput( path=output_path, - fps=runtime.metadata.fps, - output_layout=runtime.metadata.output_layout, + fps=config.fps, + output_layout=config.output_layout, enabled=environment.world_rank == 0, ) output_closed = False @@ -193,10 +284,15 @@ def close_output() -> None: resources.callback(close_output) output.open() - session = runtime.start_session(runtime.initial_input) + session = runtime.start_session(run_spec.initial_input) resources.callback(session.close) - while (request := session.next_step_request()) is not None: - output.write(session.step(runtime.prepare_step_input(request))) + for _ in range(run_spec.total_steps): + request = session.next_step_request() + if request is None: + raise RuntimeError( + "Runtime session ended before the requested MP4 step count." + ) + output.write(session.step(InferenceInput())) artifacts = tuple(output.close()) output_closed = True return artifacts diff --git a/apps/flashdreams_app/flashdreams_app/contracts.py b/apps/flashdreams_app/flashdreams_app/contracts.py index 59b8bbc34..74f952595 100644 --- a/apps/flashdreams_app/flashdreams_app/contracts.py +++ b/apps/flashdreams_app/flashdreams_app/contracts.py @@ -7,140 +7,170 @@ import argparse from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Any, Protocol, runtime_checkable +from dataclasses import dataclass +from typing import Any, Protocol, Sequence, runtime_checkable from flashdreams.infra.pipeline import StreamInferencePipelineConfig from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.runtime import InferenceInput, InferenceRuntime -from flashdreams.runtime.types import StepRequest, StepRequirements +from flashdreams.runtime import InferenceInput @dataclass(frozen=True, slots=True) -class AppConfig: - """Provider-specific CLI values normalized by the application host.""" +class AppRequest: + """Parsed invocation supplied to an application provider.""" + + mode: str + """Selected presentation mode.""" options: Mapping[str, Any] + """Parsed presentation and provider options keyed by argument destination.""" + + def __post_init__(self) -> None: + if self.mode not in ("mp4", "webrtc"): + raise ValueError(f"Unsupported application mode: {self.mode!r}.") + if not isinstance(self.options, Mapping): + raise TypeError("AppRequest.options must be a mapping.") @dataclass(frozen=True, kw_only=True, slots=True) -class RuntimeMetadata: - """Presentation facts published by an application runtime.""" +class AppConfig: + """Presentation configuration supplied by an application provider.""" model_id: str + """Stable model identity used by presentation and serving layers.""" + fps: int | float + """Output video frame rate.""" + output_layout: VideoTensorLayout + """Layout of video tensors returned by the pipeline.""" + video_width: int + """Output video width in pixels.""" + video_height: int + """Output video height in pixels.""" def __post_init__(self) -> None: if not self.model_id.strip(): - raise ValueError("RuntimeMetadata.model_id must be non-empty.") + raise ValueError("AppConfig.model_id must be non-empty.") if float(self.fps) <= 0: - raise ValueError("RuntimeMetadata.fps must be > 0.") + raise ValueError("AppConfig.fps must be > 0.") if self.video_width <= 0 or self.video_height <= 0: - raise ValueError("RuntimeMetadata video dimensions must be > 0.") - - -PipelineCacheInitializer = Callable[[Any, InferenceInput], object] + raise ValueError("AppConfig video dimensions must be > 0.") @dataclass(frozen=True, kw_only=True, slots=True) -class PipelineContract: - """Describe the model-specific operation needed to start a pipeline session. +class PipelineAppSpec: + """Mode-independent pipeline definition consumed by the application host.""" - The host already understands the standard streaming pipeline operations: - ``generate``, ``finalize``, and ``get_num_output_frames``. A provider only - supplies the cache initialization that binds its global conditioning to a - concrete pipeline implementation. - """ + pipeline_config: StreamInferencePipelineConfig + """Pipeline config that the host constructs through ``setup()``.""" - initialize_cache: PipelineCacheInitializer + initialize_cache: Callable[[Any, InferenceInput], object] + """Create one cache from the constructed pipeline and session input.""" def __post_init__(self) -> None: + if not isinstance(self.pipeline_config, StreamInferencePipelineConfig): + raise TypeError( + "PipelineAppSpec.pipeline_config must be a " + "StreamInferencePipelineConfig." + ) if not callable(self.initialize_cache): - raise TypeError("PipelineContract.initialize_cache must be callable.") + raise TypeError("PipelineAppSpec.initialize_cache must be callable.") @dataclass(frozen=True, kw_only=True, slots=True) -class PipelineAppSpec: - """Declarative application definition consumed by the generic host runtime.""" +class Mp4RunSpec: + """Finite session inputs required by the MP4 presentation path.""" - pipeline_config: StreamInferencePipelineConfig - contract: PipelineContract - metadata: RuntimeMetadata initial_input: InferenceInput + """Global conditioning used to start the finite session.""" + total_steps: int - result_metadata: Mapping[str, Any] = field(default_factory=dict) + """Number of autoregressive steps to write to the output file.""" def __post_init__(self) -> None: - if not isinstance(self.pipeline_config, StreamInferencePipelineConfig): - raise TypeError( - "PipelineAppSpec.pipeline_config must be a " - "StreamInferencePipelineConfig." - ) - if not isinstance(self.contract, PipelineContract): - raise TypeError("PipelineAppSpec.contract must be a PipelineContract.") - if not isinstance(self.metadata, RuntimeMetadata): - raise TypeError("PipelineAppSpec.metadata must be RuntimeMetadata.") if not isinstance(self.initial_input, InferenceInput): - raise TypeError("PipelineAppSpec.initial_input must be InferenceInput.") + raise TypeError("Mp4RunSpec.initial_input must be InferenceInput.") if isinstance(self.total_steps, bool) or not isinstance(self.total_steps, int): - raise TypeError("PipelineAppSpec.total_steps must be an integer.") + raise TypeError("Mp4RunSpec.total_steps must be an integer.") if self.total_steps <= 0: - raise ValueError("PipelineAppSpec.total_steps must be > 0.") - object.__setattr__( - self, - "result_metadata", - MappingProxyType(dict(self.result_metadata)), - ) + raise ValueError("Mp4RunSpec.total_steps must be > 0.") -@runtime_checkable -class AppProvider(Protocol): - """Provider module contract consumed by the application host.""" +@dataclass(frozen=True, kw_only=True, slots=True) +class WebRTCRunSpec: + """Session inputs required by the live WebRTC presentation path.""" - def add_arguments(self, parser: argparse.ArgumentParser) -> None: - """Register provider-specific arguments on the host parser.""" - ... + initial_input: InferenceInput + """Global conditioning used to start each live session.""" - def create_app_spec(self, config: AppConfig) -> PipelineAppSpec: - """Create a declarative application specification.""" - ... + def __post_init__(self) -> None: + if not isinstance(self.initial_input, InferenceInput): + raise TypeError("WebRTCRunSpec.initial_input must be InferenceInput.") -def require_pipeline_config( - config: object, *, expected_name: str | None = None -) -> StreamInferencePipelineConfig: - """Validate a pipeline provider result at the application boundary.""" - if not isinstance(config, StreamInferencePipelineConfig): - raise TypeError( - f"Pipeline provider returned {type(config).__name__}, expected " - "StreamInferencePipelineConfig." - ) - if expected_name is not None and config.name != expected_name: - raise ValueError( - f"Preset {expected_name!r} constructed pipeline {config.name!r}; " - "the preset key and pipeline name must match." - ) - return config +@dataclass(frozen=True, kw_only=True, slots=True) +class AppSpec: + """Pipeline definition paired with the selected mode's run data.""" + + config: AppConfig + """Model identity and video presentation configuration.""" + + pipeline: PipelineAppSpec + """Mode-independent pipeline definition.""" + + run: Mp4RunSpec | WebRTCRunSpec + """Inputs and limits for the selected presentation mode.""" + + def __post_init__(self) -> None: + if not isinstance(self.config, AppConfig): + raise TypeError("AppSpec.config must be AppConfig.") + if not isinstance(self.pipeline, PipelineAppSpec): + raise TypeError("AppSpec.pipeline must be PipelineAppSpec.") + if not isinstance(self.run, (Mp4RunSpec, WebRTCRunSpec)): + raise TypeError("AppSpec.run must be Mp4RunSpec or WebRTCRunSpec.") @runtime_checkable -class AppRuntime(InferenceRuntime, Protocol): - """FlashDreams runtime plus the small surface required by the app host. +class AppProvider(Protocol): + """Required interface for an installed application-provider module.""" + + def parse_options( + self, + parser: argparse.ArgumentParser, + argv: Sequence[str], + ) -> Mapping[str, Any]: + """Parse provider arguments with the selected mode's host parser. + + Args: + parser: Parser preconfigured with host-owned presentation options. + argv: Arguments remaining after the provider and mode. + + Returns: + Parsed presentation and provider options keyed by destination. + """ + ... - Sessions use the standard :class:`~flashdreams.runtime.InferenceSession` - protocol. Providers add only immutable initial inputs, per-step input - preparation, and presentation metadata. - """ + def create_app_spec(self, request: AppRequest) -> AppSpec: + """Describe the pipeline application without constructing its runtime. - metadata: RuntimeMetadata - initial_input: InferenceInput + Args: + request: Parsed invocation supplied by the application host. - def prepare_step_input( - self, request: StepRequest | StepRequirements - ) -> InferenceInput: - """Build model-facing input for one host-owned session step.""" + Returns: + Pipeline definition and run data for the selected presentation mode. + """ ... + + +__all__ = [ + "AppConfig", + "AppProvider", + "AppRequest", + "AppSpec", + "Mp4RunSpec", + "PipelineAppSpec", + "WebRTCRunSpec", +] diff --git a/apps/flashdreams_app/flashdreams_app/runtime.py b/apps/flashdreams_app/flashdreams_app/runtime.py index 20f05dff4..30c4956d7 100644 --- a/apps/flashdreams_app/flashdreams_app/runtime.py +++ b/apps/flashdreams_app/flashdreams_app/runtime.py @@ -14,11 +14,10 @@ InferenceInput, InferenceSession, StepRequest, - StepRequirements, StepResult, ) -from .contracts import PipelineAppSpec, RuntimeMetadata +from .contracts import AppConfig, PipelineAppSpec class PipelineAppRuntime: @@ -28,34 +27,28 @@ def __init__( self, *, spec: PipelineAppSpec, + config: AppConfig, device: str, ) -> None: self.pipeline = spec.pipeline_config.setup().to(device).eval() - self.metadata = spec.metadata - self.initial_input = spec.initial_input + self._config = config self._spec = spec self._closed = False - def prepare_step_input( - self, request: StepRequest | StepRequirements - ) -> InferenceInput: - """Return an empty step payload for prompt-conditioned pipelines.""" - del request - return InferenceInput() - def start_session(self, inputs: InferenceInput) -> "PipelineAppSession": - """Create a finite session with cache state isolated from other sessions.""" + """Create an open-ended session with isolated pipeline cache state.""" if self._closed: raise RuntimeError("Pipeline application runtime is closed.") return PipelineAppSession( pipeline=self.pipeline, inputs=inputs, spec=self._spec, + config=self._config, ) def peek_input_fps(self) -> float: """Return the host clock rate used for realtime presentation.""" - return float(self.metadata.fps) + return float(self._config.fps) def peek_steady_output_num_frames(self) -> int: """Return the steady-state output chunk size for presentation queues.""" @@ -74,7 +67,7 @@ def close(self) -> None: class PipelineAppSession(InferenceSession): - """Host-owned finite autoregressive session for a pipeline app spec.""" + """Host-owned autoregressive session for a pipeline application.""" def __init__( self, @@ -82,18 +75,17 @@ def __init__( pipeline: Any, inputs: InferenceInput, spec: PipelineAppSpec, + config: AppConfig, ) -> None: self._pipeline = pipeline - self._cache: object | None = spec.contract.initialize_cache(pipeline, inputs) - self._metadata = spec.metadata - self._result_metadata = spec.result_metadata - self._total_steps = spec.total_steps + self._cache: object | None = spec.initialize_cache(pipeline, inputs) + self._config = config self._step_index = 0 self._closed = False def next_step_request(self) -> StepRequest | None: - """Return the next finite rollout request, or ``None`` when complete.""" - if self._closed or self._step_index >= self._total_steps: + """Return the next rollout request, or ``None`` after session closure.""" + if self._closed: return None return StepRequest(step_index=self._step_index) @@ -102,8 +94,6 @@ def step(self, inputs: InferenceInput) -> StepResult: del inputs if self._closed: raise RuntimeError("Pipeline application session is closed.") - if self._step_index >= self._total_steps: - raise RuntimeError("Pipeline application session is complete.") if self._cache is None: raise RuntimeError("Pipeline application session has no active cache.") @@ -127,13 +117,12 @@ def step(self, inputs: InferenceInput) -> StepResult: return StepResult.from_video_chunk( step_index=index, video_chunk=video.detach(), - layout=self._metadata.output_layout, + layout=self._config.output_layout, metrics=metrics, - metadata=self._result_metadata, ) def reset(self, inputs: InferenceInput | None = None) -> None: - """Reject reset because finite sessions use isolated cache state.""" + """Reject reset because sessions use isolated cache state.""" del inputs raise RuntimeError("Create a new session instead of resetting this one.") diff --git a/apps/flashdreams_app/flashdreams_app/webrtc.py b/apps/flashdreams_app/flashdreams_app/webrtc.py index fb897c956..1690e2b53 100644 --- a/apps/flashdreams_app/flashdreams_app/webrtc.py +++ b/apps/flashdreams_app/flashdreams_app/webrtc.py @@ -6,9 +6,8 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any -from flashdreams.runtime import InferenceConfig, InferenceInput +from flashdreams.runtime import InferenceConfig, InferenceInput, InferenceRuntime from flashdreams.runtime.demo import ( DemoSpec, PreparedScenario, @@ -23,7 +22,7 @@ from flashdreams.serving.webrtc.demo import serve_webrtc_demo from flashdreams.serving.webrtc.manager import BaseWebRTCSessionManager -from .contracts import AppRuntime +from .contracts import AppConfig # TODO: Move this contract into shared FlashDreams and expand it when the @@ -45,17 +44,17 @@ class _InputProvider: deterministic_given_inputs=True, ) - def __init__(self, runtime: AppRuntime) -> None: - self._runtime = runtime + def __init__(self, initial_input: InferenceInput) -> None: + self._initial_input = initial_input def prepare_initial_input(self) -> InferenceInput: - return self._runtime.initial_input + return self._initial_input def prepare_step( self, *, request: StepRequirements, user_window: UserInputWindow ) -> PreparedStep: - del user_window - return PreparedStep(inference_input=self._runtime.prepare_step_input(request)) + del request, user_window + return PreparedStep(inference_input=InferenceInput()) def reset(self, inputs: InferenceInput | None = None) -> None: del inputs @@ -65,38 +64,55 @@ def close(self) -> None: def serve_webrtc( - *, runtime: AppRuntime, options: WebRTCOptions, device: str, world_rank: int + *, + runtime: InferenceRuntime, + config: AppConfig, + initial_input: InferenceInput, + options: WebRTCOptions, + device: str, + world_rank: int, ) -> object: - """Serve an application runtime through the shared WebRTC stack.""" - metadata = runtime.metadata + """Serve an application runtime through the shared WebRTC stack. + + Args: + runtime: Initialized application runtime. + config: Model identity and video presentation configuration. + initial_input: Global conditioning used to start live sessions. + options: WebRTC bind settings. + device: Device used by the runtime. + world_rank: Distributed rank responsible for presentation. + + Returns: + Serving backend result. + """ output = WebRTCOutputSpec( host=options.host, port=options.port, - fps=int(metadata.fps), - video_width=metadata.video_width, - video_height=metadata.video_height, - preload_name=metadata.model_id, + fps=int(config.fps), + video_width=config.video_width, + video_height=config.video_height, + preload_name=config.model_id, ) spec = DemoSpec( - model_id=metadata.model_id, + model_id=config.model_id, input_mode="webrtc", output=output, - config=InferenceConfig(model_id=metadata.model_id, device=device), + config=InferenceConfig(model_id=config.model_id, device=device), ) - scenario = PreparedScenario(initial_inputs=runtime.initial_input) + scenario = PreparedScenario(initial_inputs=initial_input) def create_model_input_provider( spec: DemoSpec, scenario: PreparedScenario, ) -> _InputProvider: del spec, scenario - return _InputProvider(runtime) + return _InputProvider(initial_input) manager = BaseWebRTCSessionManager( runtime=runtime, runtime_config=output, - fps=int(metadata.fps), - identity=metadata.model_id, + fps=int(config.fps), + identity=config.model_id, shared_host=RuntimeHost(runtime), shared_spec=spec, shared_scenario=scenario, @@ -105,8 +121,8 @@ def create_model_input_provider( ) return serve_webrtc_demo( output=output, - model_id=metadata.model_id, + model_id=config.model_id, session_manager=manager, - app_resources=WebRTCAppResources(preload_name=metadata.model_id), + app_resources=WebRTCAppResources(preload_name=config.model_id), world_rank=world_rank, ) diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py index 0c5aee0f1..94e7f1584 100644 --- a/apps/flashdreams_app/tests/test_cli.py +++ b/apps/flashdreams_app/tests/test_cli.py @@ -4,16 +4,21 @@ from __future__ import annotations import argparse +from collections.abc import Mapping, Sequence from types import ModuleType, SimpleNamespace from typing import Any, cast +import flashdreams_app import pytest import torch from flashdreams_app import ( + AppConfig, AppProvider, + AppRequest, + AppSpec, + Mp4RunSpec, PipelineAppSpec, - PipelineContract, - RuntimeMetadata, + WebRTCRunSpec, cli, ) @@ -23,6 +28,20 @@ pytestmark = pytest.mark.ci_cpu +def test_public_package_surface_contains_only_provider_contracts() -> None: + assert flashdreams_app.__all__ == [ + "AppConfig", + "AppProvider", + "AppRequest", + "AppSpec", + "Mp4RunSpec", + "PipelineAppSpec", + "WebRTCRunSpec", + ] + assert "initial_input" not in PipelineAppSpec.__dataclass_fields__ + assert "total_steps" not in PipelineAppSpec.__dataclass_fields__ + + def test_host_drives_runtime_api_and_owns_file_output( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -77,28 +96,37 @@ def initialize_cache(pipeline: object, inputs: InferenceInput) -> object: provider = ModuleType("fake_app") - def add_arguments(parser: argparse.ArgumentParser) -> None: - del parser - calls.append("provider.add_arguments") - - setattr(provider, "add_arguments", add_arguments) - setattr( - provider, - "create_app_spec", - lambda config: PipelineAppSpec( - pipeline_config=pipeline_config, - contract=PipelineContract(initialize_cache=initialize_cache), - metadata=RuntimeMetadata( + def parse_options( + parser: argparse.ArgumentParser, argv: Sequence[str] + ) -> Mapping[str, object]: + calls.append("provider.parse_options") + parser.add_argument("--model-option", required=True) + return vars(parser.parse_args(argv)) + + def create_app_spec(request: AppRequest) -> AppSpec: + assert request.mode == "mp4" + assert request.options["model_option"] == "enabled" + calls.append("provider.create_app_spec") + return AppSpec( + config=AppConfig( model_id="fake", fps=24, output_layout="tchw", video_width=64, video_height=64, ), - initial_input=InferenceInput(global_conditioning={"prompt": "test"}), - total_steps=1, - ), - ) + pipeline=PipelineAppSpec( + pipeline_config=pipeline_config, + initialize_cache=initialize_cache, + ), + run=Mp4RunSpec( + initial_input=InferenceInput(global_conditioning={"prompt": "test"}), + total_steps=1, + ), + ) + + setattr(provider, "parse_options", parse_options) + setattr(provider, "create_app_spec", create_app_spec) monkeypatch.setattr(cli, "load_provider", lambda _: provider) class Output: @@ -116,9 +144,21 @@ def close(self) -> tuple[object, ...]: return () monkeypatch.setattr(cli, "FileOutput", Output) - cli.run(["fake-app", "mp4", "--device", "cpu", "--output", "result.mp4"]) + cli.run( + [ + "fake-app", + "mp4", + "--device", + "cpu", + "--output", + "result.mp4", + "--model-option", + "enabled", + ] + ) assert calls == [ - "provider.add_arguments", + "provider.parse_options", + "provider.create_app_spec", "pipeline.init", "pipeline.to", "pipeline.eval", @@ -134,13 +174,101 @@ def close(self) -> tuple[object, ...]: ] +def test_run_delegates_options_and_dispatches_webrtc( + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = ModuleType("fake_app") + pipeline_config = StreamInferencePipelineConfig( + _target=cast(Any, object), + name="fake", + diffusion_model=cast(Any, None), + ) + initial_input = InferenceInput(global_conditioning={"prompt": "test"}) + + def parse_options( + parser: argparse.ArgumentParser, argv: Sequence[str] + ) -> Mapping[str, object]: + assert tuple(argv) == ( + "--host", + "127.0.0.1", + "--port", + "9000", + "--prompt", + "test", + ) + parser.add_argument("--prompt", required=True) + return vars(parser.parse_args(argv)) + + def create_app_spec(request: AppRequest) -> AppSpec: + assert request.mode == "webrtc" + assert request.options["prompt"] == "test" + return AppSpec( + config=AppConfig( + model_id="fake", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + ), + pipeline=PipelineAppSpec( + pipeline_config=pipeline_config, + initialize_cache=lambda pipeline, inputs: object(), + ), + run=WebRTCRunSpec(initial_input=initial_input), + ) + + setattr(provider, "parse_options", parse_options) + setattr(provider, "create_app_spec", create_app_spec) + monkeypatch.setattr(cli, "load_provider", lambda _: provider) + environment = cli._Environment(device="cpu", world_rank=0, world_size=1) + monkeypatch.setattr(cli, "_initialize_environment", lambda device: environment) + runtime = object() + monkeypatch.setattr(cli, "PipelineAppRuntime", lambda **kwargs: runtime) + + captured: dict[str, object] = {} + + def run_webrtc(**kwargs: object) -> tuple[object, ...]: + captured.update(kwargs) + return () + + monkeypatch.setattr(cli, "_run_webrtc", run_webrtc) + monkeypatch.setattr( + cli, + "_run_mp4", + lambda **kwargs: pytest.fail("WebRTC mode must not launch the MP4 path."), + ) + + assert ( + cli.run( + [ + "fake-app", + "webrtc", + "--host", + "127.0.0.1", + "--port", + "9000", + "--prompt", + "test", + ] + ) + == () + ) + assert captured["runtime"] is runtime + assert captured["run_spec"] == WebRTCRunSpec(initial_input=initial_input) + assert captured["environment"] is environment + args = captured["args"] + assert isinstance(args, argparse.Namespace) + assert args.host == "127.0.0.1" + assert args.port == 9000 + + def test_app_provider_protocol_requires_both_methods() -> None: provider = ModuleType("provider") - setattr(provider, "add_arguments", lambda parser: None) + setattr(provider, "parse_options", lambda parser, argv: {}) setattr(provider, "create_app_spec", lambda config: None) assert isinstance(provider, AppProvider) - delattr(provider, "add_arguments") + delattr(provider, "parse_options") assert not isinstance(provider, AppProvider) @@ -162,21 +290,30 @@ def test_load_provider_rejects_module_outside_contract( def test_host_exposes_only_supported_output_modes() -> None: - mode_action = next( - action for action in cli.build_parser()._actions if action.dest == "mode" - ) - assert mode_action.choices == ("mp4", "webrtc") + route = cli._parse_provider_and_mode(["fake-app", "webrtc", "--prompt", "x"]) + assert route.provider == "fake-app" + assert route.mode == "webrtc" + assert route.remaining_argv == ("--prompt", "x") + + with pytest.raises(SystemExit): + cli._parse_provider_and_mode(["fake-app", "unsupported"]) def test_host_does_not_expose_pipeline_execution_options() -> None: - destinations = {action.dest for action in cli.build_parser()._actions} - assert "compile" not in destinations - assert "cuda_graph" not in destinations + for mode in ("mp4", "webrtc"): + destinations = { + action.dest for action in cli.build_parser("fake-app", mode)._actions + } + assert "compile" not in destinations + assert "cuda_graph" not in destinations def test_host_exposes_only_minimal_webrtc_options() -> None: - destinations = {action.dest for action in cli.build_parser()._actions} + destinations = { + action.dest for action in cli.build_parser("fake-app", "webrtc")._actions + } assert {"host", "port"} <= destinations + assert "output" not in destinations assert { "warmup_chunks", "warmup_timeout_s", @@ -187,6 +324,14 @@ def test_host_exposes_only_minimal_webrtc_options() -> None: }.isdisjoint(destinations) +def test_mp4_parser_exposes_only_file_presentation_options() -> None: + destinations = { + action.dest for action in cli.build_parser("fake-app", "mp4")._actions + } + assert {"device", "output"} <= destinations + assert {"host", "port"}.isdisjoint(destinations) + + def test_webrtc_path_owns_serving_options_and_runtime_close( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -194,19 +339,6 @@ def test_webrtc_path_owns_serving_options_and_runtime_close( captured: dict[str, object] = {} class Runtime: - metadata = RuntimeMetadata( - model_id="fake", - fps=24, - output_layout="tchw", - video_width=64, - video_height=64, - ) - initial_input = InferenceInput() - - def prepare_step_input(self, request: object) -> InferenceInput: - del request - return InferenceInput() - def start_session(self, inputs: InferenceInput) -> Any: del inputs raise AssertionError("The WebRTC path must not start an MP4 session.") @@ -221,6 +353,14 @@ def serve(**kwargs: object) -> None: monkeypatch.setattr(cli, "serve_webrtc", serve) result = cli._run_webrtc( runtime=Runtime(), + config=AppConfig( + model_id="fake", + fps=24, + output_layout="tchw", + video_width=64, + video_height=64, + ), + run_spec=WebRTCRunSpec(initial_input=InferenceInput()), args=argparse.Namespace( host="127.0.0.1", port=9000, @@ -236,3 +376,4 @@ def serve(**kwargs: object) -> None: assert options.host == "127.0.0.1" assert options.port == 9000 assert captured["device"] == "cpu" + assert isinstance(captured["initial_input"], InferenceInput) diff --git a/apps/flashdreams_app/tests/test_webrtc.py b/apps/flashdreams_app/tests/test_webrtc.py index 6be395dbe..c1e4b66c3 100644 --- a/apps/flashdreams_app/tests/test_webrtc.py +++ b/apps/flashdreams_app/tests/test_webrtc.py @@ -4,7 +4,7 @@ from __future__ import annotations import pytest -from flashdreams_app import RuntimeMetadata, webrtc +from flashdreams_app import AppConfig, webrtc from flashdreams.runtime import InferenceInput, StepRequest, StepResult @@ -26,19 +26,6 @@ def close(self) -> None: class _Runtime: - metadata = RuntimeMetadata( - model_id="fake-app", - fps=16, - output_layout="tchw", - video_width=96, - video_height=64, - ) - initial_input = InferenceInput() - - def prepare_step_input(self, request: object) -> InferenceInput: - del request - return InferenceInput() - def start_session(self, inputs: InferenceInput) -> _Session: del inputs return _Session() @@ -58,8 +45,17 @@ def fake_serve(**kwargs: object) -> str: monkeypatch.setattr(webrtc, "serve_webrtc_demo", fake_serve) runtime = _Runtime() + initial_input = InferenceInput() result = webrtc.serve_webrtc( runtime=runtime, + config=AppConfig( + model_id="fake-app", + fps=16, + output_layout="tchw", + video_width=96, + video_height=64, + ), + initial_input=initial_input, options=webrtc.WebRTCOptions( host="127.0.0.1", port=8080, @@ -82,5 +78,5 @@ def fake_serve(**kwargs: object) -> str: assert session_manager._shared_host.runtime is runtime assert session_manager.runtime_config is output assert session_manager._shared_scenario is not None - assert session_manager._shared_scenario.initial_inputs is runtime.initial_input + assert session_manager._shared_scenario.initial_inputs is initial_input assert callable(session_manager._shared_model_input_provider_factory) diff --git a/apps/t2v_app/README.md b/apps/t2v_app/README.md index c17ce47a1..b4b49abb6 100644 --- a/apps/t2v_app/README.md +++ b/apps/t2v_app/README.md @@ -1,9 +1,10 @@ # T2V App Provider `t2v-app` is a model provider for the generic `flashdreams-app` host. It -returns a declarative `PipelineAppSpec`; it does not implement a runtime or -session and does not own setup, stepping, finalization, cleanup, MP4 writing, -or WebRTC. +returns a declarative `AppSpec` containing a mode-independent +`PipelineAppSpec`, presentation `AppConfig`, and mode-specific run data; it +does not implement a runtime or session and does not own setup, stepping, +finalization, cleanup, MP4 writing, or WebRTC. The provider loads a YAML preset catalog through `flashdreams.core.pipeline_presets` and asks the selected pipeline provider to @@ -29,8 +30,10 @@ uv run flashdreams-app t2v-app mp4 \ --output outputs/waterfall.mp4 ``` -Every YAML preset must specify `provider`, all six runtime/presentation fields, -and the provider-owned `pipeline` options. FlashDreams' +Every YAML preset must specify `provider`, the common runtime/presentation +fields, and the provider-owned `pipeline` options. `total_blocks` is optional; +MP4 mode requires either that preset default or `--total-blocks`, while WebRTC +does not use a finite step count. FlashDreams' `ObjectGraphPipelineProvider` supports these declarative nodes: - `_target: module:attribute` imports and calls a config class with the other @@ -45,6 +48,11 @@ and reference it from `provider`. Preset YAML is trusted configuration because provider and object-graph references import Python objects. At the provider boundary, T2V contributes only its preset selection, -conditioning values, presentation metadata, and a cache initializer that maps -the prompt and pixel dimensions to the selected pipeline. `flashdreams-app` -constructs and drives the resulting pipeline. +conditioning values, presentation configuration, and a cache initializer that +maps the prompt and pixel dimensions to the selected pipeline. Session +conditioning and finite MP4 length live in mode-specific run specs; +`flashdreams-app` constructs and drives the resulting pipeline. + +See the [`flashdreams-app` provider contract](../flashdreams_app/README.md#provider-contract) +for the required entry points, pipeline and run-spec fields, and a minimal +provider implementation. diff --git a/apps/t2v_app/pyproject.toml b/apps/t2v_app/pyproject.toml index 0e0b61081..745ca85f2 100644 --- a/apps/t2v_app/pyproject.toml +++ b/apps/t2v_app/pyproject.toml @@ -17,6 +17,11 @@ dependencies = ["flashdreams", "flashdreams-app"] flashdreams = { workspace = true } flashdreams-app = { workspace = true } +[tool.pyright] +extraPaths = ["../../flashdreams", "../flashdreams_app"] +venvPath = "../.." +venv = ".venv" + [tool.setuptools.packages.find] where = ["."] diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py index 9a8680c96..26a6bad0e 100644 --- a/apps/t2v_app/t2v_app/__init__.py +++ b/apps/t2v_app/t2v_app/__init__.py @@ -3,6 +3,6 @@ """Text-to-video pipeline application for ``flashdreams-app``.""" -from .provider import add_arguments, create_app_spec +from .provider import create_app_spec, parse_options -__all__ = ["add_arguments", "create_app_spec"] +__all__ = ["create_app_spec", "parse_options"] diff --git a/apps/t2v_app/t2v_app/presets.py b/apps/t2v_app/t2v_app/presets.py index 60b5e30f3..2c9e8ea1e 100644 --- a/apps/t2v_app/t2v_app/presets.py +++ b/apps/t2v_app/t2v_app/presets.py @@ -31,14 +31,14 @@ ) from flashdreams.infra.postprocess import VideoTensorLayout -_RUNTIME_FIELDS = { +_REQUIRED_RUNTIME_FIELDS = { "prompt", - "total_blocks", "pixel_height", "pixel_width", "fps", "output_layout", } +_OPTIONAL_RUNTIME_FIELDS = {"total_blocks"} _VIDEO_LAYOUTS = {"tchw", "btchw", "bcthw", "bvtchw"} @@ -49,9 +49,6 @@ class RuntimePresetOptions: prompt: str """Default text prompt.""" - total_blocks: int - """Number of autoregressive chunks in one finite session.""" - pixel_height: int """Output video height in pixels.""" @@ -64,6 +61,9 @@ class RuntimePresetOptions: 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, @@ -92,7 +92,12 @@ def load_preset_catalog( def _load_runtime_options(value: object, *, path: str) -> RuntimePresetOptions: runtime = _mapping(value, path=path) - _require_exact_fields(runtime, expected=_RUNTIME_FIELDS, 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)) @@ -101,15 +106,17 @@ def _load_runtime_options(value: object, *, path: str) -> RuntimePresetOptions: ) return RuntimePresetOptions( prompt=_nonempty_string(runtime["prompt"], path=f"{path}.prompt"), - total_blocks=_positive_int( - runtime["total_blocks"], path=f"{path}.total_blocks" - ), 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") + ), ) @@ -122,12 +129,16 @@ def _mapping(value: object, *, path: str) -> Mapping[str, object]: return cast(dict[str, object], mapping) -def _require_exact_fields( - value: Mapping[str, object], *, expected: set[str], path: str +def _require_fields( + value: Mapping[str, object], + *, + required: set[str], + optional: set[str], + path: str, ) -> None: fields = {str(key) for key in value} - missing = expected - fields - unknown = fields - expected + missing = required - fields + unknown = fields - required - optional if missing or unknown: details: list[str] = [] if missing: diff --git a/apps/t2v_app/t2v_app/provider.py b/apps/t2v_app/t2v_app/provider.py index 753004cd1..692ff9157 100644 --- a/apps/t2v_app/t2v_app/provider.py +++ b/apps/t2v_app/t2v_app/provider.py @@ -6,16 +6,17 @@ from __future__ import annotations import argparse -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Any from flashdreams_app import ( AppConfig, + AppRequest, + AppSpec, + Mp4RunSpec, PipelineAppSpec, - PipelineContract, - RuntimeMetadata, - require_pipeline_config, + WebRTCRunSpec, ) from flashdreams.core.pipeline_presets import load_pipeline_provider @@ -32,8 +33,19 @@ FIELD_FPS = "fps" -def add_arguments(parser: argparse.ArgumentParser) -> None: - """Register T2V conditioning and rollout options on the host parser.""" +def parse_options( + parser: argparse.ArgumentParser, + argv: Sequence[str], +) -> Mapping[str, Any]: + """Parse T2V options with the selected mode's host parser. + + Args: + parser: Parser preconfigured with host-owned presentation options. + argv: Arguments remaining after the provider and mode. + + Returns: + Parsed presentation and T2V options keyed by destination. + """ parser.add_argument( "--preset-config", type=Path, @@ -48,15 +60,30 @@ def add_arguments(parser: argparse.ArgumentParser) -> None: 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) + return vars(parser.parse_args(argv)) + +def create_app_spec(request: AppRequest) -> AppSpec: + """Describe a T2V pipeline application without constructing its runtime. -def create_app_spec(config: AppConfig) -> PipelineAppSpec: - """Describe a T2V pipeline application without constructing its runtime.""" - options = config.options + Args: + request: Parsed host and T2V command-line values. + + Returns: + Pipeline selection, initial conditioning, and presentation data. + """ + options = request.options preset_id, preset = _resolve_preset(options) scenario = _scenario(options, preset.runtime) pipeline_config = _create_pipeline_config(preset_id, preset) - return _build_app_spec(pipeline_config, scenario, preset.runtime) + return AppSpec( + config=_build_app_config(scenario, preset.runtime), + pipeline=PipelineAppSpec( + pipeline_config=pipeline_config, + initialize_cache=_initialize_cache, + ), + run=_build_run_spec(request.mode, options, scenario, preset.runtime), + ) def _resolve_preset( @@ -73,43 +100,71 @@ def _create_pipeline_config( ) -> StreamInferencePipelineConfig: """Create the pipeline config selected by a resolved preset.""" pipeline_provider = load_pipeline_provider(preset.provider) - return require_pipeline_config( - pipeline_provider.create_pipeline_config( - preset_id=preset_id, - options=preset.pipeline, - ), - expected_name=preset_id, + 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 _build_app_spec( - pipeline_config: StreamInferencePipelineConfig, +def _build_app_config( scenario: Mapping[str, object], runtime_options: RuntimePresetOptions, -) -> PipelineAppSpec: - """Build the host-facing application spec from resolved T2V data.""" - return PipelineAppSpec( - pipeline_config=pipeline_config, - contract=PipelineContract(initialize_cache=_initialize_cache), - metadata=RuntimeMetadata( - model_id="t2v-app", - fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), - output_layout=runtime_options.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 - ), - ), - initial_input=InferenceInput(global_conditioning=scenario), - total_steps=_required_int( - scenario[FIELD_TOTAL_BLOCKS], name=FIELD_TOTAL_BLOCKS +) -> AppConfig: + """Build presentation configuration from the resolved T2V scenario.""" + return AppConfig( + model_id="t2v-app", + fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), + output_layout=runtime_options.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 ), - result_metadata={FIELD_PROMPT: scenario[FIELD_PROMPT]}, ) +def _build_run_spec( + mode: str, + options: Mapping[str, object], + scenario: Mapping[str, object], + defaults: RuntimePresetOptions, +) -> Mp4RunSpec | WebRTCRunSpec: + """Build the selected presentation mode's session data.""" + initial_input = InferenceInput( + global_conditioning={ + FIELD_PROMPT: scenario[FIELD_PROMPT], + FIELD_PIXEL_HEIGHT: scenario[FIELD_PIXEL_HEIGHT], + FIELD_PIXEL_WIDTH: scenario[FIELD_PIXEL_WIDTH], + } + ) + if mode == "webrtc": + return WebRTCRunSpec(initial_input=initial_input) + if mode == "mp4": + total_steps = options.get(FIELD_TOTAL_BLOCKS) + if total_steps is None: + total_steps = defaults.total_blocks + if total_steps is None: + raise ValueError( + "MP4 mode requires --total-blocks or runtime.total_blocks in " + "the selected preset." + ) + return Mp4RunSpec( + initial_input=initial_input, + total_steps=_required_int(total_steps, name=FIELD_TOTAL_BLOCKS), + ) + raise ValueError(f"Unsupported presentation mode: {mode!r}.") + + def _initialize_cache(pipeline: Any, inputs: InferenceInput) -> object: """Bind T2V prompt and dimensions to a new pipeline cache.""" scenario = _scenario_from_inputs(inputs) @@ -140,9 +195,6 @@ def _scenario( prompt = _resolve_prompt(defaults.prompt if prompt_value is None else prompt_value) scenario = { FIELD_PROMPT: prompt, - FIELD_TOTAL_BLOCKS: _option_or_default( - options, FIELD_TOTAL_BLOCKS, defaults.total_blocks - ), FIELD_PIXEL_HEIGHT: _option_or_default( options, FIELD_PIXEL_HEIGHT, defaults.pixel_height ), @@ -152,7 +204,6 @@ def _scenario( FIELD_FPS: _option_or_default(options, FIELD_FPS, defaults.fps), } for name in ( - FIELD_TOTAL_BLOCKS, FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH, FIELD_FPS, @@ -166,10 +217,8 @@ def _scenario_from_inputs(inputs: InferenceInput) -> Mapping[str, object]: source = inputs.global_conditioning required = ( FIELD_PROMPT, - FIELD_TOTAL_BLOCKS, FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH, - FIELD_FPS, ) missing = tuple(name for name in required if name not in source) if missing: @@ -177,10 +226,8 @@ def _scenario_from_inputs(inputs: InferenceInput) -> Mapping[str, object]: scenario = dict(source) scenario[FIELD_PROMPT] = _resolve_prompt(scenario[FIELD_PROMPT]) for name in ( - FIELD_TOTAL_BLOCKS, FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH, - FIELD_FPS, ): if _required_int(scenario[name], name=name) <= 0: raise ValueError(f"{name} must be > 0.") @@ -225,4 +272,4 @@ def _optional_string(value: object) -> str | None: return value.strip() -__all__ = ["add_arguments", "create_app_spec"] +__all__ = ["create_app_spec", "parse_options"] diff --git a/apps/t2v_app/tests/test_pipeline_provider.py b/apps/t2v_app/tests/test_pipeline_provider.py index a069bffe4..28e17600a 100644 --- a/apps/t2v_app/tests/test_pipeline_provider.py +++ b/apps/t2v_app/tests/test_pipeline_provider.py @@ -20,11 +20,10 @@ from pathlib import Path import pytest -from flashdreams_app import require_pipeline_config +from t2v_app import provider as app_provider from t2v_app.presets import load_preset_catalog from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict -from flashdreams.core.pipeline_presets import load_pipeline_provider from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig from flashdreams.recipes.wan import Wan21TransformerConfig @@ -35,14 +34,7 @@ def test_packaged_yaml_constructs_default_pipeline_config() -> None: catalog = load_preset_catalog() preset_id, preset = catalog.resolve(None) - provider = load_pipeline_provider(preset.provider) - config = require_pipeline_config( - provider.create_pipeline_config( - preset_id=preset_id, - options=preset.pipeline, - ), - expected_name=preset_id, - ) + config = app_provider._create_pipeline_config(preset_id, preset) assert preset_id == "causal-forcing-wan2.1-t2v-1.3b-chunkwise" assert config.name == preset_id @@ -81,6 +73,31 @@ def test_catalog_rejects_incomplete_runtime_options(tmp_path: Path) -> None: 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() diff --git a/apps/t2v_app/tests/test_provider.py b/apps/t2v_app/tests/test_provider.py index 78c3d307e..fcc0ba8eb 100644 --- a/apps/t2v_app/tests/test_provider.py +++ b/apps/t2v_app/tests/test_provider.py @@ -10,7 +10,14 @@ import pytest import t2v_app -from flashdreams_app import AppConfig, AppProvider, PipelineAppSpec +from flashdreams_app import ( + AppProvider, + AppRequest, + AppSpec, + Mp4RunSpec, + PipelineAppSpec, + WebRTCRunSpec, +) from t2v_app import provider from t2v_app.presets import ( PipelinePreset, @@ -27,15 +34,14 @@ def test_provider_module_conforms_to_host_contract() -> None: assert isinstance(t2v_app, AppProvider) -def test_t2v_provider_registers_model_options() -> None: +def test_t2v_provider_parses_model_options() -> None: parser = argparse.ArgumentParser() - provider.add_arguments(parser) - args = parser.parse_args(["--prompt", "A waterfall"]) - assert args.prompt == "A waterfall" - assert args.preset_id is None - assert args.preset_config is None - assert not hasattr(args, "backend") - assert not hasattr(args, "compile") + options = provider.parse_options(parser, ["--prompt", "A waterfall"]) + assert options["prompt"] == "A waterfall" + assert options["preset_id"] is None + assert options["preset_config"] is None + assert "backend" not in options + assert "compile" not in options def test_create_app_spec_returns_data_without_constructing_pipeline( @@ -84,7 +90,8 @@ def create_pipeline_config( monkeypatch.setattr(provider, "load_pipeline_provider", lambda _: Provider()) created = provider.create_app_spec( - AppConfig( + AppRequest( + mode="mp4", options={ "preset_config": None, "preset_id": None, @@ -93,14 +100,42 @@ def create_pipeline_config( "pixel_height": None, "pixel_width": None, "fps": None, - } + }, ) ) - assert isinstance(created, PipelineAppSpec) - assert created.pipeline_config is pipeline_config - assert created.initial_input.global_conditioning["prompt"] == "A waterfall" - assert created.metadata.video_width == 96 - assert created.metadata.fps == 12 - assert created.total_steps == 2 + assert isinstance(created, AppSpec) + assert isinstance(created.pipeline, PipelineAppSpec) + assert created.pipeline.pipeline_config is pipeline_config + assert created.config.video_width == 96 + assert created.config.fps == 12 + assert isinstance(created.run, Mp4RunSpec) + assert created.run.initial_input.global_conditioning["prompt"] == "A waterfall" + assert created.run.total_steps == 2 assert not pipeline_constructed + + +def test_webrtc_run_spec_does_not_require_total_steps() -> None: + defaults = RuntimePresetOptions( + prompt="default prompt", + pixel_height=64, + pixel_width=96, + fps=12, + output_layout="tchw", + ) + scenario = { + provider.FIELD_PROMPT: "A waterfall", + provider.FIELD_PIXEL_HEIGHT: 64, + provider.FIELD_PIXEL_WIDTH: 96, + provider.FIELD_FPS: 12, + } + + run_spec = provider._build_run_spec( + "webrtc", + {"total_blocks": None}, + scenario, + defaults, + ) + + assert isinstance(run_spec, WebRTCRunSpec) + assert run_spec.initial_input.global_conditioning["prompt"] == "A waterfall" From e412b8903f63b9e6c11219205aeaeb24b5091c42 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 05:44:54 +0000 Subject: [PATCH 08/10] Refactor FlashDreams application runtime ABI Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 165 ++++---- apps/flashdreams_app/README.md | 114 ------ .../flashdreams_app/__init__.py | 24 -- apps/flashdreams_app/flashdreams_app/cli.py | 327 --------------- .../flashdreams_app/contracts.py | 176 -------- .../flashdreams_app/outputs.py | 44 -- .../flashdreams_app/runtime.py | 158 -------- apps/flashdreams_app/tests/test_cli.py | 379 ------------------ apps/t2v_app/README.md | 76 ++-- apps/t2v_app/pyproject.toml | 8 +- apps/t2v_app/t2v_app/__init__.py | 6 +- apps/t2v_app/t2v_app/application.py | 199 +++++++++ apps/t2v_app/t2v_app/provider.py | 275 ------------- apps/t2v_app/t2v_app/runtime.py | 93 +++++ apps/t2v_app/t2v_app/session.py | 189 +++++++++ apps/t2v_app/tests/test_application.py | 253 ++++++++++++ apps/t2v_app/tests/test_pipeline_provider.py | 4 +- apps/t2v_app/tests/test_provider.py | 141 ------- flashdreams_runner/README.md | 152 +++++++ flashdreams_runner/__init__.py | 28 ++ flashdreams_runner/cli.py | 240 +++++++++++ flashdreams_runner/contracts.py | 260 ++++++++++++ flashdreams_runner/modes.py | 172 ++++++++ flashdreams_runner/outputs.py | 111 +++++ .../pyproject.toml | 11 +- flashdreams_runner/tests/test_cli.py | 306 ++++++++++++++ .../tests/test_webrtc.py | 65 ++- .../webrtc.py | 89 ++-- pyproject.toml | 3 +- uv.lock | 28 +- 30 files changed, 2251 insertions(+), 1845 deletions(-) delete mode 100644 apps/flashdreams_app/README.md delete mode 100644 apps/flashdreams_app/flashdreams_app/__init__.py delete mode 100644 apps/flashdreams_app/flashdreams_app/cli.py delete mode 100644 apps/flashdreams_app/flashdreams_app/contracts.py delete mode 100644 apps/flashdreams_app/flashdreams_app/outputs.py delete mode 100644 apps/flashdreams_app/flashdreams_app/runtime.py delete mode 100644 apps/flashdreams_app/tests/test_cli.py create mode 100644 apps/t2v_app/t2v_app/application.py delete mode 100644 apps/t2v_app/t2v_app/provider.py create mode 100644 apps/t2v_app/t2v_app/runtime.py create mode 100644 apps/t2v_app/t2v_app/session.py create mode 100644 apps/t2v_app/tests/test_application.py delete mode 100644 apps/t2v_app/tests/test_provider.py create mode 100644 flashdreams_runner/README.md create mode 100644 flashdreams_runner/__init__.py create mode 100644 flashdreams_runner/cli.py create mode 100644 flashdreams_runner/contracts.py create mode 100644 flashdreams_runner/modes.py create mode 100644 flashdreams_runner/outputs.py rename {apps/flashdreams_app => flashdreams_runner}/pyproject.toml (61%) create mode 100644 flashdreams_runner/tests/test_cli.py rename {apps/flashdreams_app => flashdreams_runner}/tests/test_webrtc.py (54%) rename {apps/flashdreams_app/flashdreams_app => flashdreams_runner}/webrtc.py (57%) diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md index f39886847..51d4273d3 100644 --- a/PR_DESCRIPTION.md +++ b/PR_DESCRIPTION.md @@ -1,41 +1,37 @@ ## Summary -Add a model-neutral `flashdreams-app` entrypoint and a declarative `t2v-app` -provider for running FlashDreams text-to-video pipelines as MP4 jobs or WebRTC -services. +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 host owns process initialization, pipeline construction, the -runtime/session lifecycle, autoregressive stepping, finalization, cleanup, and -presentation. The T2V package only selects a pipeline preset and supplies the -model-specific conditioning and cache-initialization callback. +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-app t2v-app {mp4 | webrtc} +uv run flashdreams-runner t2v-app {mp4 | replay | webrtc | none} | v -+----------------------+ request spec +----------------------+ -| flashdreams_app | --------------------> | t2v_app | -| generic host | <-------------------- | declarative provider | -+----------+-----------+ AppSpec +----------+-----------+ - | | - | constructs and drives | reads - v v -+----------------------+ +------------------------+ -| FlashDreams runtime | | pipeline_presets.yaml | -| API (black box) | | + PipelineProvider | -+----------+-----------+ +------------------------+ - | - | video chunks - v - +----+----+ - | | - v v -+----------+ +----------+ -| MP4 file | | WebRTC | -| artifact | | stream | -+----------+ +----------+ ++--------------------------+ 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.generate | StepResult +--------------------+ +| present output | ++------------+-------------+ + | + +-----+-----+----------------+ + | | | + v v v + MP4/replay WebRTC None ``` ## Entrypoint examples @@ -43,7 +39,16 @@ uv run flashdreams-app t2v-app {mp4 | webrtc} Generate an MP4 with the packaged default preset: ```bash -uv run flashdreams-app t2v-app mp4 \ +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 ``` @@ -51,78 +56,78 @@ uv run flashdreams-app t2v-app mp4 \ Serve the same application through WebRTC: ```bash -uv run flashdreams-app t2v-app webrtc \ +uv run flashdreams-runner t2v-app webrtc \ --prompt "A waterfall" ``` -Select a packaged preset explicitly: +Run without presentation or artifacts: ```bash -uv run flashdreams-app t2v-app mp4 \ - --preset-id self-forcing-wan2.1-t2v-1.3b \ - --prompt "A neon-lit city at night" \ - --output outputs/city.mp4 +uv run flashdreams-runner t2v-app none \ + --steps 2 \ + --prompt "A waterfall" ``` -Load a custom YAML preset catalog: +Select a packaged preset explicitly: ```bash -uv run flashdreams-app t2v-app webrtc \ - --preset-config /path/to/pipeline-presets.yaml \ - --preset-id my-t2v-preset \ - --prompt "A waterfall at sunset" +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`. -## Architecture - -- Add the `flashdreams-app` workspace package and console entrypoint. -- Define a minimal provider boundary: - `create_app_spec(AppRequest) -> AppSpec`. -- Keep the provider surface data-first: a mode-independent pipeline spec plus - an `AppConfig` for presentation and mode-specific MP4 or WebRTC run data. -- Require provider modules to conform to `AppProvider` with - `parse_options(parser, argv)` and `create_app_spec(request)`. -- Add the host-owned `PipelineAppRuntime` and `PipelineAppSession`. -- Keep pipeline setup, `generate`/`finalize`, step tracking, cache release, and - runtime closure in the host. -- Add host-owned MP4 and WebRTC presentation paths without a runtime adapter. -- Type both presentation paths directly against the shared `InferenceRuntime` - contract. -- Parse only the provider and presentation mode in the host, then let the - provider extend the mode-specific parser and parse all remaining arguments - through `parse_options(parser, argv)`. -- Keep pipeline-specific execution behavior encapsulated by the selected - pipeline config and its `setup()` implementation. - -## T2V provider and presets - -- Add the `t2v-app` workspace package. -- Describe T2V through an `AppSpec` rather than implementing another runtime - or session. -- Add packaged YAML presets for causal-forcing and self-forcing WAN pipelines. -- Resolve pipeline providers directly from the YAML catalog without depending - on the runner registry. -- Support trusted declarative `_target`, `_ref`, and `_tuple` nodes for pipeline - object graphs. +## 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.generate()` 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.generate()`. +- 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 the causal-forcing and self-forcing configurations. -- Let the shared asynchronous demo and WebRTC manager consume runtime objects - directly when no model adapter is needed. + across causal-forcing and self-forcing configurations. ## Validation -- `57` affected CPU tests pass. -- Full `pre-commit run -a` passes, including Ruff formatting, lockfile - validation, and `ty` type checking. -- `flashdreams-app t2v-app mp4 --help` resolves the installed provider and - composes host and T2V arguments correctly. +- 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/flashdreams_app/README.md b/apps/flashdreams_app/README.md deleted file mode 100644 index 27ab31da7..000000000 --- a/apps/flashdreams_app/README.md +++ /dev/null @@ -1,114 +0,0 @@ -# FlashDreams App Host - -`flashdreams-app` is the model-neutral application host. It finds a provider -distribution in the active Python environment, asks it for a declarative -pipeline application spec, constructs the runtime, creates a session, runs the -session loop, and writes presentation artifacts itself. - -```bash -uv run flashdreams-app t2v-app mp4 --output o.mp4 --prompt "A waterfall" -uv run flashdreams-app t2v-app webrtc --prompt "A waterfall" -``` - -## Provider contract - -A compatible package must expose an importable module with exactly two public -entry points: - -- `parse_options(parser, argv)`, which extends the selected mode's host parser - with provider-specific flags, parses the remaining arguments, and returns the - parsed values as a mapping. -- `create_app_spec(request: flashdreams_app.AppRequest)`, returning a - mode-aware `AppSpec` containing a shared `PipelineAppSpec` and either an - `Mp4RunSpec` or `WebRTCRunSpec`. - -The module structurally conforms to `flashdreams_app.AppProvider`; the host -validates this contract when it loads the installed provider. - -`PipelineAppSpec` contains only mode-independent runtime data: - -| Provider-supplied field | Purpose | -|---|---| -| `pipeline_config` | A `StreamInferencePipelineConfig`; the host calls `setup()` and owns the resulting pipeline. | -| `initialize_cache` | A callback receiving the constructed pipeline and session input; it returns the session cache. | - -`AppSpec.config` is an `AppConfig` containing the model identity, video -dimensions, frame rate, and output tensor layout used by host-owned -presentation. `AppRequest` contains only the selected mode and parsed invocation -options passed into the provider. - -Invocation data belongs to the selected presentation mode: - -| Run spec | Required data | -|---|---| -| `Mp4RunSpec` | `initial_input` and the finite `total_steps` written to the file. | -| `WebRTCRunSpec` | `initial_input` for each live session; no finite step count. | - -The following is a complete provider skeleton: - -```python -import argparse -from collections.abc import Mapping, Sequence -from typing import Any - -from flashdreams_app import ( - AppConfig, - AppRequest, - AppSpec, - Mp4RunSpec, - PipelineAppSpec, - WebRTCRunSpec, -) -from flashdreams.runtime import InferenceInput -from my_model.config import MY_PIPELINE_CONFIG - - -def parse_options( - parser: argparse.ArgumentParser, - argv: Sequence[str], -) -> Mapping[str, Any]: - parser.add_argument("--prompt", required=True) - return vars(parser.parse_args(argv)) - - -def create_app_spec(request: AppRequest) -> AppSpec: - prompt = str(request.options["prompt"]) - initial_input = InferenceInput(global_conditioning={"prompt": prompt}) - if request.mode == "mp4": - run = Mp4RunSpec(initial_input=initial_input, total_steps=4) - else: - run = WebRTCRunSpec(initial_input=initial_input) - return AppSpec( - config=AppConfig( - model_id="my-model", - fps=24, - output_layout="tchw", - video_width=832, - video_height=480, - ), - pipeline=PipelineAppSpec( - pipeline_config=MY_PIPELINE_CONFIG, - initialize_cache=_initialize_cache, - ), - run=run, - ) - - -def _initialize_cache(pipeline: Any, inputs: InferenceInput) -> object: - prompt = str(inputs.global_conditioning["prompt"]) - return pipeline.initialize_cache( - text=[prompt], - image=None, - height=60, - width=104, - ) -``` - -The host first parses only the provider and presentation mode. It then gives the -provider a parser containing that mode's host-owned options and all remaining -arguments. The host owns process/distributed initialization, pipeline setup, -runtime/session lifecycle, `generate`/`finalize` stepping, MP4 writing, and -WebRTC serving. Providers only parse their options, select/configure a pipeline, -and describe how global conditioning initializes its cache. Pipeline-specific -execution options, including compilation and CUDA graphs, remain encapsulated -by the pipeline config and its `setup()` implementation. diff --git a/apps/flashdreams_app/flashdreams_app/__init__.py b/apps/flashdreams_app/flashdreams_app/__init__.py deleted file mode 100644 index fc79e7875..000000000 --- a/apps/flashdreams_app/flashdreams_app/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Public contract for FlashDreams application providers.""" - -from .contracts import ( - AppConfig, - AppProvider, - AppRequest, - AppSpec, - Mp4RunSpec, - PipelineAppSpec, - WebRTCRunSpec, -) - -__all__ = [ - "AppConfig", - "AppProvider", - "AppRequest", - "AppSpec", - "Mp4RunSpec", - "PipelineAppSpec", - "WebRTCRunSpec", -] diff --git a/apps/flashdreams_app/flashdreams_app/cli.py b/apps/flashdreams_app/flashdreams_app/cli.py deleted file mode 100644 index 59d42b9ee..000000000 --- a/apps/flashdreams_app/flashdreams_app/cli.py +++ /dev/null @@ -1,327 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Command line host for independently installed FlashDreams app providers.""" - -from __future__ import annotations - -import argparse -import importlib -from contextlib import ExitStack -from dataclasses import dataclass -from importlib import metadata -from pathlib import Path -from typing import Sequence - -import torch - -from flashdreams.runtime import InferenceInput, InferenceRuntime -from flashdreams.runtime.demo.bootstrap import ( - configure_logging, - initialize_cuda_distributed, -) -from flashdreams.runtime.output import OutputArtifact - -from .contracts import ( - AppConfig, - AppProvider, - AppRequest, - AppSpec, - Mp4RunSpec, - WebRTCRunSpec, -) -from .outputs import FileOutput -from .runtime import PipelineAppRuntime -from .webrtc import WebRTCOptions, serve_webrtc - - -@dataclass(frozen=True, slots=True) -class _ProviderAndMode: - """Top-level route parsed before loading an application provider.""" - - provider: str - """Installed provider distribution name.""" - - mode: str - """Selected presentation mode.""" - - remaining_argv: tuple[str, ...] - """Arguments delegated to the provider parser.""" - - -def build_parser(provider: str, mode: str) -> argparse.ArgumentParser: - """Build the selected mode's host-owned options parser. - - Args: - provider: Provider name displayed in command usage. - mode: Selected presentation mode. - - Returns: - Parser ready for the provider to extend and invoke. - """ - parser = argparse.ArgumentParser(prog=f"flashdreams-app {provider} {mode}") - parser.add_argument("--device", default="cuda", help="Runtime device") - if mode == "mp4": - parser.add_argument("--output", type=Path, required=True, help="MP4 path") - elif 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") - else: - raise ValueError(f"Unsupported application mode: {mode!r}.") - return parser - - -def load_provider(distribution_name: str) -> AppProvider: - """Load an installed provider that satisfies the host contract.""" - try: - distribution = metadata.distribution(distribution_name) - except metadata.PackageNotFoundError as exc: - raise ValueError( - f"Provider 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, AppProvider): - return module - incompatible.append(candidate) - if incompatible: - names = ", ".join(repr(name) for name in incompatible) - raise TypeError( - f"Provider distribution {distribution_name!r} exposes module(s) " - f"{names}, but none satisfy AppProvider. Providers must define " - "parse_options(parser, argv) and create_app_spec(request)." - ) - raise ValueError( - f"Provider {distribution_name!r} does not expose an importable Python module." - ) - - -def run(argv: Sequence[str] | None = None) -> tuple[OutputArtifact, ...]: - """Dispatch one provider session to its selected presentation path. - - Args: - argv: Command-line arguments; ``None`` reads the process arguments. - - Returns: - Artifacts produced by the selected path. WebRTC returns an empty tuple. - """ - route = _parse_provider_and_mode(argv) - provider = load_provider(route.provider) - options = provider.parse_options( - build_parser(route.provider, route.mode), - route.remaining_argv, - ) - app_spec = _require_app_spec( - provider.create_app_spec(AppRequest(mode=route.mode, options=options)), - provider_name=route.provider, - mode=route.mode, - ) - args = argparse.Namespace(**options) - environment = _initialize_environment(args.device) - runtime: InferenceRuntime = PipelineAppRuntime( - spec=app_spec.pipeline, - config=app_spec.config, - device=environment.device, - ) - return _launch_mode( - mode=route.mode, - runtime=runtime, - app_spec=app_spec, - args=args, - environment=environment, - ) - - -def _parse_provider_and_mode( - argv: Sequence[str] | None, -) -> _ProviderAndMode: - """Parse the provider and mode while preserving all remaining arguments.""" - parser = argparse.ArgumentParser(prog="flashdreams-app", add_help=False) - parser.add_argument("provider", help="Installed provider distribution") - parser.add_argument("mode", choices=("mp4", "webrtc")) - args, remaining_argv = parser.parse_known_args(argv) - return _ProviderAndMode( - provider=args.provider, - mode=args.mode, - remaining_argv=tuple(remaining_argv), - ) - - -def _require_app_spec( - value: object, - *, - provider_name: str, - mode: str, -) -> AppSpec: - """Validate a provider result before constructing its runtime.""" - if not isinstance(value, AppSpec): - raise TypeError( - f"Provider {provider_name!r} create_app_spec() returned " - f"{type(value).__name__}, expected AppSpec." - ) - if mode == "webrtc" and not isinstance(value.run, WebRTCRunSpec): - raise TypeError("WebRTC mode requires WebRTCRunSpec from the provider.") - if mode == "mp4" and not isinstance(value.run, Mp4RunSpec): - raise TypeError("MP4 mode requires Mp4RunSpec from the provider.") - return value - - -def _launch_mode( - *, - mode: str, - runtime: InferenceRuntime, - app_spec: AppSpec, - args: argparse.Namespace, - environment: "_Environment", -) -> tuple[OutputArtifact, ...]: - """Launch the host-owned presentation path selected by the user.""" - if mode == "webrtc": - assert isinstance(app_spec.run, WebRTCRunSpec) - return _run_webrtc( - runtime=runtime, - config=app_spec.config, - run_spec=app_spec.run, - args=args, - environment=environment, - ) - if mode == "mp4": - assert isinstance(app_spec.run, Mp4RunSpec) - return _run_mp4( - runtime=runtime, - config=app_spec.config, - run_spec=app_spec.run, - output_path=args.output, - environment=environment, - ) - raise AssertionError(f"Unsupported presentation mode: {mode!r}.") - - -def _run_webrtc( - *, - runtime: InferenceRuntime, - config: AppConfig, - run_spec: WebRTCRunSpec, - args: argparse.Namespace, - environment: "_Environment", -) -> tuple[OutputArtifact, ...]: - """Run the WebRTC serving path and close its runtime. - - Args: - runtime: Initialized application runtime. - config: Model identity and video presentation configuration. - run_spec: Initial conditioning for live sessions. - args: Parsed host and WebRTC arguments. - environment: Initialized process and distributed environment. - - Returns: - An empty tuple because WebRTC does not create file artifacts. - """ - try: - serve_webrtc( - runtime=runtime, - config=config, - initial_input=run_spec.initial_input, - options=WebRTCOptions( - host=args.host, - port=args.port, - ), - device=environment.device, - world_rank=environment.world_rank, - ) - finally: - runtime.close() - return () - - -def _run_mp4( - *, - runtime: InferenceRuntime, - config: AppConfig, - run_spec: Mp4RunSpec, - output_path: Path, - environment: "_Environment", -) -> tuple[OutputArtifact, ...]: - """Run the finite MP4 generation path and close all owned resources. - - Args: - runtime: Initialized application runtime. - config: Model identity and video presentation configuration. - run_spec: Initial conditioning and finite generation length. - output_path: Destination for the generated MP4. - environment: Initialized process and distributed environment. - - Returns: - Artifacts emitted by the file output target. - """ - - with ExitStack() as resources: - resources.callback(runtime.close) - output = FileOutput( - path=output_path, - fps=config.fps, - output_layout=config.output_layout, - enabled=environment.world_rank == 0, - ) - output_closed = False - - def close_output() -> None: - if not output_closed: - output.close() - - resources.callback(close_output) - output.open() - session = runtime.start_session(run_spec.initial_input) - resources.callback(session.close) - for _ in range(run_spec.total_steps): - request = session.next_step_request() - if request is None: - raise RuntimeError( - "Runtime session ended before the requested MP4 step count." - ) - output.write(session.step(InferenceInput())) - artifacts = tuple(output.close()) - output_closed = True - return artifacts - - -class _Environment: - def __init__(self, *, device: str, world_rank: int, world_size: int) -> None: - self.device = device - self.world_rank = world_rank - self.world_size = world_size - - -def _initialize_environment(device: str) -> _Environment: - """Initialize logging, CUDA placement, and distributed state in the host.""" - 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: - """Console-script entry point.""" - artifacts = run() - # File modes return persistent artifacts whose URI is the output location; - # live modes such as WebRTC return no artifacts. - for artifact in artifacts: - print(artifact.uri) diff --git a/apps/flashdreams_app/flashdreams_app/contracts.py b/apps/flashdreams_app/flashdreams_app/contracts.py deleted file mode 100644 index 74f952595..000000000 --- a/apps/flashdreams_app/flashdreams_app/contracts.py +++ /dev/null @@ -1,176 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Data-first provider boundary used by :mod:`flashdreams_app`.""" - -from __future__ import annotations - -import argparse -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from typing import Any, Protocol, Sequence, runtime_checkable - -from flashdreams.infra.pipeline import StreamInferencePipelineConfig -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.runtime import InferenceInput - - -@dataclass(frozen=True, slots=True) -class AppRequest: - """Parsed invocation supplied to an application provider.""" - - mode: str - """Selected presentation mode.""" - - options: Mapping[str, Any] - """Parsed presentation and provider options keyed by argument destination.""" - - def __post_init__(self) -> None: - if self.mode not in ("mp4", "webrtc"): - raise ValueError(f"Unsupported application mode: {self.mode!r}.") - if not isinstance(self.options, Mapping): - raise TypeError("AppRequest.options must be a mapping.") - - -@dataclass(frozen=True, kw_only=True, slots=True) -class AppConfig: - """Presentation configuration supplied by an application provider.""" - - model_id: str - """Stable model identity used by presentation and serving layers.""" - - fps: int | float - """Output video frame rate.""" - - output_layout: VideoTensorLayout - """Layout of video tensors returned by the pipeline.""" - - video_width: int - """Output video width in pixels.""" - - video_height: int - """Output video height in pixels.""" - - 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.") - - -@dataclass(frozen=True, kw_only=True, slots=True) -class PipelineAppSpec: - """Mode-independent pipeline definition consumed by the application host.""" - - pipeline_config: StreamInferencePipelineConfig - """Pipeline config that the host constructs through ``setup()``.""" - - initialize_cache: Callable[[Any, InferenceInput], object] - """Create one cache from the constructed pipeline and session input.""" - - def __post_init__(self) -> None: - if not isinstance(self.pipeline_config, StreamInferencePipelineConfig): - raise TypeError( - "PipelineAppSpec.pipeline_config must be a " - "StreamInferencePipelineConfig." - ) - if not callable(self.initialize_cache): - raise TypeError("PipelineAppSpec.initialize_cache must be callable.") - - -@dataclass(frozen=True, kw_only=True, slots=True) -class Mp4RunSpec: - """Finite session inputs required by the MP4 presentation path.""" - - initial_input: InferenceInput - """Global conditioning used to start the finite session.""" - - total_steps: int - """Number of autoregressive steps to write to the output file.""" - - def __post_init__(self) -> None: - if not isinstance(self.initial_input, InferenceInput): - raise TypeError("Mp4RunSpec.initial_input must be InferenceInput.") - if isinstance(self.total_steps, bool) or not isinstance(self.total_steps, int): - raise TypeError("Mp4RunSpec.total_steps must be an integer.") - if self.total_steps <= 0: - raise ValueError("Mp4RunSpec.total_steps must be > 0.") - - -@dataclass(frozen=True, kw_only=True, slots=True) -class WebRTCRunSpec: - """Session inputs required by the live WebRTC presentation path.""" - - initial_input: InferenceInput - """Global conditioning used to start each live session.""" - - def __post_init__(self) -> None: - if not isinstance(self.initial_input, InferenceInput): - raise TypeError("WebRTCRunSpec.initial_input must be InferenceInput.") - - -@dataclass(frozen=True, kw_only=True, slots=True) -class AppSpec: - """Pipeline definition paired with the selected mode's run data.""" - - config: AppConfig - """Model identity and video presentation configuration.""" - - pipeline: PipelineAppSpec - """Mode-independent pipeline definition.""" - - run: Mp4RunSpec | WebRTCRunSpec - """Inputs and limits for the selected presentation mode.""" - - def __post_init__(self) -> None: - if not isinstance(self.config, AppConfig): - raise TypeError("AppSpec.config must be AppConfig.") - if not isinstance(self.pipeline, PipelineAppSpec): - raise TypeError("AppSpec.pipeline must be PipelineAppSpec.") - if not isinstance(self.run, (Mp4RunSpec, WebRTCRunSpec)): - raise TypeError("AppSpec.run must be Mp4RunSpec or WebRTCRunSpec.") - - -@runtime_checkable -class AppProvider(Protocol): - """Required interface for an installed application-provider module.""" - - def parse_options( - self, - parser: argparse.ArgumentParser, - argv: Sequence[str], - ) -> Mapping[str, Any]: - """Parse provider arguments with the selected mode's host parser. - - Args: - parser: Parser preconfigured with host-owned presentation options. - argv: Arguments remaining after the provider and mode. - - Returns: - Parsed presentation and provider options keyed by destination. - """ - ... - - def create_app_spec(self, request: AppRequest) -> AppSpec: - """Describe the pipeline application without constructing its runtime. - - Args: - request: Parsed invocation supplied by the application host. - - Returns: - Pipeline definition and run data for the selected presentation mode. - """ - ... - - -__all__ = [ - "AppConfig", - "AppProvider", - "AppRequest", - "AppSpec", - "Mp4RunSpec", - "PipelineAppSpec", - "WebRTCRunSpec", -] diff --git a/apps/flashdreams_app/flashdreams_app/outputs.py b/apps/flashdreams_app/flashdreams_app/outputs.py deleted file mode 100644 index d6439b45e..000000000 --- a/apps/flashdreams_app/flashdreams_app/outputs.py +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Host-owned presentation targets.""" - -from __future__ import annotations - -from pathlib import Path - -from flashdreams.infra.postprocess import VideoTensorLayout -from flashdreams.runtime.output import OutputArtifact -from flashdreams.runtime.types import StepResult -from flashdreams.runtime.video_output import Mp4VideoOutputTarget - - -class FileOutput: - """Collect generated chunks and write one MP4 file when the run completes.""" - - def __init__( - self, - *, - path: Path, - fps: int | float, - output_layout: VideoTensorLayout, - enabled: bool = True, - ) -> None: - self._target = Mp4VideoOutputTarget( - output_path=path, - fps=fps, - output_layout=output_layout, - enabled=enabled, - ) - - def open(self) -> None: - """Open the underlying video writer.""" - self._target.open() - - def write(self, result: StepResult) -> None: - """Append a generated chunk.""" - self._target.write(result) - - def close(self) -> tuple[OutputArtifact, ...]: - """Finalize the MP4 and return its artifact metadata.""" - return tuple(self._target.close()) diff --git a/apps/flashdreams_app/flashdreams_app/runtime.py b/apps/flashdreams_app/flashdreams_app/runtime.py deleted file mode 100644 index 30c4956d7..000000000 --- a/apps/flashdreams_app/flashdreams_app/runtime.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Host-owned runtime and session for streaming pipeline applications.""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import Any - -import torch - -from flashdreams.runtime import ( - InferenceInput, - InferenceSession, - StepRequest, - StepResult, -) - -from .contracts import AppConfig, PipelineAppSpec - - -class PipelineAppRuntime: - """Instantiate and own the reusable pipeline behind an application spec.""" - - def __init__( - self, - *, - spec: PipelineAppSpec, - config: AppConfig, - device: str, - ) -> None: - self.pipeline = spec.pipeline_config.setup().to(device).eval() - self._config = config - self._spec = spec - self._closed = False - - def start_session(self, inputs: InferenceInput) -> "PipelineAppSession": - """Create an open-ended session with isolated pipeline cache state.""" - if self._closed: - raise RuntimeError("Pipeline application runtime is closed.") - return PipelineAppSession( - pipeline=self.pipeline, - inputs=inputs, - spec=self._spec, - config=self._config, - ) - - def peek_input_fps(self) -> float: - """Return the host clock rate used for realtime presentation.""" - return float(self._config.fps) - - def peek_steady_output_num_frames(self) -> int: - """Return the steady-state output chunk size for presentation queues.""" - return int(self.pipeline.get_num_output_frames(1)) - - def close(self) -> None: - """Release the shared pipeline and accelerator allocator state.""" - if self._closed: - return - self._closed = True - close = getattr(self.pipeline, "close", None) - if callable(close): - close() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - -class PipelineAppSession(InferenceSession): - """Host-owned autoregressive session for a pipeline application.""" - - def __init__( - self, - *, - pipeline: Any, - inputs: InferenceInput, - spec: PipelineAppSpec, - config: AppConfig, - ) -> None: - self._pipeline = pipeline - self._cache: object | None = spec.initialize_cache(pipeline, inputs) - self._config = config - self._step_index = 0 - self._closed = False - - def next_step_request(self) -> StepRequest | None: - """Return the next rollout request, or ``None`` after session closure.""" - if self._closed: - return None - return StepRequest(step_index=self._step_index) - - def step(self, inputs: InferenceInput) -> StepResult: - """Generate and finalize one autoregressive video chunk.""" - del inputs - if self._closed: - raise RuntimeError("Pipeline application session is closed.") - if self._cache is None: - raise RuntimeError("Pipeline application session has no active cache.") - - index = self._step_index - video = self._pipeline.generate( - autoregressive_index=index, - cache=self._cache, - ) - metrics = _metrics( - self._pipeline.finalize( - autoregressive_index=index, - cache=self._cache, - ) - ) - self._step_index += 1 - if not isinstance(video, torch.Tensor): - raise TypeError( - "Pipeline generate() must return a torch.Tensor, got " - f"{type(video).__name__}." - ) - return StepResult.from_video_chunk( - step_index=index, - video_chunk=video.detach(), - layout=self._config.output_layout, - metrics=metrics, - ) - - def reset(self, inputs: InferenceInput | None = None) -> None: - """Reject reset because sessions use isolated cache state.""" - del inputs - raise RuntimeError("Create a new session instead of resetting this one.") - - def close(self) -> None: - """Release session-local cache state.""" - if self._closed: - return - self._closed = True - cache = self._cache - self._cache = None - close = getattr(cache, "close", None) - if callable(close): - close() - - -def _metrics(value: object) -> Mapping[str, float | int]: - if value is None: - return {} - if not isinstance(value, Mapping): - raise TypeError( - "Pipeline finalize() must return a metrics mapping or None, got " - f"{type(value).__name__}." - ) - metrics = dict(value) - invalid = tuple( - key for key, metric in metrics.items() if not isinstance(metric, (int, float)) - ) - if invalid: - raise TypeError(f"Pipeline metrics must be numeric; invalid keys: {invalid}.") - return metrics - - -__all__ = ["PipelineAppRuntime", "PipelineAppSession"] diff --git a/apps/flashdreams_app/tests/test_cli.py b/apps/flashdreams_app/tests/test_cli.py deleted file mode 100644 index 94e7f1584..000000000 --- a/apps/flashdreams_app/tests/test_cli.py +++ /dev/null @@ -1,379 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import argparse -from collections.abc import Mapping, Sequence -from types import ModuleType, SimpleNamespace -from typing import Any, cast - -import flashdreams_app -import pytest -import torch -from flashdreams_app import ( - AppConfig, - AppProvider, - AppRequest, - AppSpec, - Mp4RunSpec, - PipelineAppSpec, - WebRTCRunSpec, - cli, -) - -from flashdreams.infra.pipeline import StreamInferencePipelineConfig -from flashdreams.runtime import InferenceInput, StepResult - -pytestmark = pytest.mark.ci_cpu - - -def test_public_package_surface_contains_only_provider_contracts() -> None: - assert flashdreams_app.__all__ == [ - "AppConfig", - "AppProvider", - "AppRequest", - "AppSpec", - "Mp4RunSpec", - "PipelineAppSpec", - "WebRTCRunSpec", - ] - assert "initial_input" not in PipelineAppSpec.__dataclass_fields__ - assert "total_steps" not in PipelineAppSpec.__dataclass_fields__ - - -def test_host_drives_runtime_api_and_owns_file_output( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[str] = [] - - class Cache: - def close(self) -> None: - calls.append("cache.close") - - class Pipeline: - def __init__(self, config: object) -> None: - assert config is pipeline_config - calls.append("pipeline.init") - - def to(self, device: str) -> "Pipeline": - assert device == "cpu" - calls.append("pipeline.to") - return self - - def eval(self) -> "Pipeline": - calls.append("pipeline.eval") - return self - - def generate(self, *, autoregressive_index: int, cache: object) -> torch.Tensor: - assert autoregressive_index == 0 - assert isinstance(cache, Cache) - calls.append("pipeline.generate") - return torch.zeros((1, 3, 2, 2)) - - def finalize( - self, *, autoregressive_index: int, cache: object - ) -> dict[str, float]: - assert autoregressive_index == 0 - assert isinstance(cache, Cache) - calls.append("pipeline.finalize") - return {"step_ms": 1.0} - - def close(self) -> None: - calls.append("pipeline.close") - - pipeline_config = StreamInferencePipelineConfig( - _target=cast(Any, Pipeline), - name="fake", - diffusion_model=cast(Any, None), - ) - - def initialize_cache(pipeline: object, inputs: InferenceInput) -> object: - assert isinstance(pipeline, Pipeline) - assert inputs.global_conditioning["prompt"] == "test" - calls.append("contract.initialize_cache") - return Cache() - - provider = ModuleType("fake_app") - - def parse_options( - parser: argparse.ArgumentParser, argv: Sequence[str] - ) -> Mapping[str, object]: - calls.append("provider.parse_options") - parser.add_argument("--model-option", required=True) - return vars(parser.parse_args(argv)) - - def create_app_spec(request: AppRequest) -> AppSpec: - assert request.mode == "mp4" - assert request.options["model_option"] == "enabled" - calls.append("provider.create_app_spec") - return AppSpec( - config=AppConfig( - model_id="fake", - fps=24, - output_layout="tchw", - video_width=64, - video_height=64, - ), - pipeline=PipelineAppSpec( - pipeline_config=pipeline_config, - initialize_cache=initialize_cache, - ), - run=Mp4RunSpec( - initial_input=InferenceInput(global_conditioning={"prompt": "test"}), - total_steps=1, - ), - ) - - setattr(provider, "parse_options", parse_options) - setattr(provider, "create_app_spec", create_app_spec) - monkeypatch.setattr(cli, "load_provider", lambda _: provider) - - class Output: - def __init__(self, **_: object) -> None: - calls.append("output.init") - - def open(self) -> None: - calls.append("output.open") - - def write(self, result: StepResult) -> None: - calls.append("output.write") - - def close(self) -> tuple[object, ...]: - calls.append("output.close") - return () - - monkeypatch.setattr(cli, "FileOutput", Output) - cli.run( - [ - "fake-app", - "mp4", - "--device", - "cpu", - "--output", - "result.mp4", - "--model-option", - "enabled", - ] - ) - assert calls == [ - "provider.parse_options", - "provider.create_app_spec", - "pipeline.init", - "pipeline.to", - "pipeline.eval", - "output.init", - "output.open", - "contract.initialize_cache", - "pipeline.generate", - "pipeline.finalize", - "output.write", - "output.close", - "cache.close", - "pipeline.close", - ] - - -def test_run_delegates_options_and_dispatches_webrtc( - monkeypatch: pytest.MonkeyPatch, -) -> None: - provider = ModuleType("fake_app") - pipeline_config = StreamInferencePipelineConfig( - _target=cast(Any, object), - name="fake", - diffusion_model=cast(Any, None), - ) - initial_input = InferenceInput(global_conditioning={"prompt": "test"}) - - def parse_options( - parser: argparse.ArgumentParser, argv: Sequence[str] - ) -> Mapping[str, object]: - assert tuple(argv) == ( - "--host", - "127.0.0.1", - "--port", - "9000", - "--prompt", - "test", - ) - parser.add_argument("--prompt", required=True) - return vars(parser.parse_args(argv)) - - def create_app_spec(request: AppRequest) -> AppSpec: - assert request.mode == "webrtc" - assert request.options["prompt"] == "test" - return AppSpec( - config=AppConfig( - model_id="fake", - fps=24, - output_layout="tchw", - video_width=64, - video_height=64, - ), - pipeline=PipelineAppSpec( - pipeline_config=pipeline_config, - initialize_cache=lambda pipeline, inputs: object(), - ), - run=WebRTCRunSpec(initial_input=initial_input), - ) - - setattr(provider, "parse_options", parse_options) - setattr(provider, "create_app_spec", create_app_spec) - monkeypatch.setattr(cli, "load_provider", lambda _: provider) - environment = cli._Environment(device="cpu", world_rank=0, world_size=1) - monkeypatch.setattr(cli, "_initialize_environment", lambda device: environment) - runtime = object() - monkeypatch.setattr(cli, "PipelineAppRuntime", lambda **kwargs: runtime) - - captured: dict[str, object] = {} - - def run_webrtc(**kwargs: object) -> tuple[object, ...]: - captured.update(kwargs) - return () - - monkeypatch.setattr(cli, "_run_webrtc", run_webrtc) - monkeypatch.setattr( - cli, - "_run_mp4", - lambda **kwargs: pytest.fail("WebRTC mode must not launch the MP4 path."), - ) - - assert ( - cli.run( - [ - "fake-app", - "webrtc", - "--host", - "127.0.0.1", - "--port", - "9000", - "--prompt", - "test", - ] - ) - == () - ) - assert captured["runtime"] is runtime - assert captured["run_spec"] == WebRTCRunSpec(initial_input=initial_input) - assert captured["environment"] is environment - args = captured["args"] - assert isinstance(args, argparse.Namespace) - assert args.host == "127.0.0.1" - assert args.port == 9000 - - -def test_app_provider_protocol_requires_both_methods() -> None: - provider = ModuleType("provider") - setattr(provider, "parse_options", lambda parser, argv: {}) - setattr(provider, "create_app_spec", lambda config: None) - assert isinstance(provider, AppProvider) - - delattr(provider, "parse_options") - assert not isinstance(provider, AppProvider) - - -def test_load_provider_rejects_module_outside_contract( - monkeypatch: pytest.MonkeyPatch, -) -> None: - module = ModuleType("invalid_provider") - distribution = SimpleNamespace(metadata={"Name": "invalid-app"}) - monkeypatch.setattr(cli.metadata, "distribution", lambda name: distribution) - monkeypatch.setattr( - cli.metadata, - "packages_distributions", - lambda: {"invalid_provider": ["invalid-app"]}, - ) - monkeypatch.setattr(cli.importlib, "import_module", lambda name: module) - - with pytest.raises(TypeError, match="none satisfy AppProvider"): - cli.load_provider("invalid-app") - - -def test_host_exposes_only_supported_output_modes() -> None: - route = cli._parse_provider_and_mode(["fake-app", "webrtc", "--prompt", "x"]) - assert route.provider == "fake-app" - assert route.mode == "webrtc" - assert route.remaining_argv == ("--prompt", "x") - - with pytest.raises(SystemExit): - cli._parse_provider_and_mode(["fake-app", "unsupported"]) - - -def test_host_does_not_expose_pipeline_execution_options() -> None: - for mode in ("mp4", "webrtc"): - destinations = { - action.dest for action in cli.build_parser("fake-app", mode)._actions - } - assert "compile" not in destinations - assert "cuda_graph" not in destinations - - -def test_host_exposes_only_minimal_webrtc_options() -> None: - destinations = { - action.dest for action in cli.build_parser("fake-app", "webrtc")._actions - } - assert {"host", "port"} <= destinations - assert "output" not in destinations - assert { - "warmup_chunks", - "warmup_timeout_s", - "client_liveness_timeout_s", - "encoder_backend", - "encoder_bitrate_bps", - "encoder_gop", - }.isdisjoint(destinations) - - -def test_mp4_parser_exposes_only_file_presentation_options() -> None: - destinations = { - action.dest for action in cli.build_parser("fake-app", "mp4")._actions - } - assert {"device", "output"} <= destinations - assert {"host", "port"}.isdisjoint(destinations) - - -def test_webrtc_path_owns_serving_options_and_runtime_close( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls: list[str] = [] - captured: dict[str, object] = {} - - class Runtime: - def start_session(self, inputs: InferenceInput) -> Any: - del inputs - raise AssertionError("The WebRTC path must not start an MP4 session.") - - def close(self) -> None: - calls.append("runtime.close") - - def serve(**kwargs: object) -> None: - calls.append("serve_webrtc") - captured.update(kwargs) - - monkeypatch.setattr(cli, "serve_webrtc", serve) - result = cli._run_webrtc( - runtime=Runtime(), - config=AppConfig( - model_id="fake", - fps=24, - output_layout="tchw", - video_width=64, - video_height=64, - ), - run_spec=WebRTCRunSpec(initial_input=InferenceInput()), - args=argparse.Namespace( - host="127.0.0.1", - port=9000, - ), - environment=cli._Environment(device="cpu", world_rank=0, world_size=1), - ) - - assert result == () - assert calls == ["serve_webrtc", "runtime.close"] - assert captured["world_rank"] == 0 - options = captured["options"] - assert isinstance(options, cli.WebRTCOptions) - assert options.host == "127.0.0.1" - assert options.port == 9000 - assert captured["device"] == "cpu" - assert isinstance(captured["initial_input"], InferenceInput) diff --git a/apps/t2v_app/README.md b/apps/t2v_app/README.md index b4b49abb6..7145e356a 100644 --- a/apps/t2v_app/README.md +++ b/apps/t2v_app/README.md @@ -1,40 +1,50 @@ -# T2V App Provider +# T2V Example Application -`t2v-app` is a model provider for the generic `flashdreams-app` host. It -returns a declarative `AppSpec` containing a mode-independent -`PipelineAppSpec`, presentation `AppConfig`, and mode-specific run data; it -does not implement a runtime or session and does not own setup, stepping, -finalization, cleanup, MP4 writing, or WebRTC. +`t2v-app` is an example implementation of the `flashdreams-runner` +application ABI. Its public module exports only `create_runtime(arguments)`. -The provider loads a YAML preset catalog through -`flashdreams.core.pipeline_presets` and asks the selected pipeline provider to -construct a FlashDreams `StreamInferencePipelineConfig`. It does not use the -`flashdreams.runner_configs` registry. The packaged catalog is -[`t2v_app/pipeline_presets.yaml`](t2v_app/pipeline_presets.yaml); pass -`--preset-config` to use another catalog. +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-app t2v-app mp4 \ +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-app t2v-app webrtc \ +uv run flashdreams-runner t2v-app webrtc \ --preset-id self-forcing-wan2.1-t2v-1.3b \ --prompt "A neon-lit city at night" -uv run flashdreams-app t2v-app mp4 \ - --preset-config /path/to/presets.yaml \ - --preset-id my-t2v-preset \ - --prompt "A waterfall" \ - --output outputs/waterfall.mp4 +uv run flashdreams-runner t2v-app none \ + --steps 2 \ + --prompt "A waterfall" ``` -Every YAML preset must specify `provider`, the common runtime/presentation -fields, and the provider-owned `pipeline` options. `total_blocks` is optional; -MP4 mode requires either that preset default or `--total-blocks`, while WebRTC -does not use a finite step count. FlashDreams' -`ObjectGraphPipelineProvider` supports these declarative nodes: +## 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. WebRTC does not impose a finite +step count. + +FlashDreams' `ObjectGraphPipelineProvider` supports these trusted declarative +nodes: - `_target: module:attribute` imports and calls a config class with the other mapping entries as keyword arguments. @@ -42,17 +52,9 @@ does not use a finite step count. FlashDreams' without calling it. - `_tuple: [...]` preserves tuple-valued config fields. -A custom package can expose a zero-argument provider class (or provider -instance) implementing `flashdreams.core.pipeline_presets.PipelineProvider` -and reference it from `provider`. Preset YAML is trusted configuration because -provider and object-graph references import Python objects. - -At the provider boundary, T2V contributes only its preset selection, -conditioning values, presentation configuration, and a cache initializer that -maps the prompt and pixel dimensions to the selected pipeline. Session -conditioning and finite MP4 length live in mode-specific run specs; -`flashdreams-app` constructs and drives the resulting pipeline. +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-app` provider contract](../flashdreams_app/README.md#provider-contract) -for the required entry points, pipeline and run-spec fields, and a minimal -provider implementation. +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 index 745ca85f2..7694afd3e 100644 --- a/apps/t2v_app/pyproject.toml +++ b/apps/t2v_app/pyproject.toml @@ -8,17 +8,17 @@ build-backend = "setuptools.build_meta" [project] name = "t2v-app" version = "0.1.0" -description = "Text-to-video pipeline provider for flashdreams-app" +description = "Text-to-video example application for flashdreams-runner" readme = "README.md" requires-python = ">=3.10" -dependencies = ["flashdreams", "flashdreams-app"] +dependencies = ["flashdreams", "flashdreams-runner"] [tool.uv.sources] flashdreams = { workspace = true } -flashdreams-app = { workspace = true } +flashdreams-runner = { workspace = true } [tool.pyright] -extraPaths = ["../../flashdreams", "../flashdreams_app"] +extraPaths = ["../..", "../../flashdreams"] venvPath = "../.." venv = ".venv" diff --git a/apps/t2v_app/t2v_app/__init__.py b/apps/t2v_app/t2v_app/__init__.py index 26a6bad0e..90a014158 100644 --- a/apps/t2v_app/t2v_app/__init__.py +++ b/apps/t2v_app/t2v_app/__init__.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Text-to-video pipeline application for ``flashdreams-app``.""" +"""Text-to-video example application for ``flashdreams-runner``.""" -from .provider import create_app_spec, parse_options +from .application import create_runtime -__all__ = ["create_app_spec", "parse_options"] +__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..a747efea7 --- /dev/null +++ b/apps/t2v_app/t2v_app/application.py @@ -0,0 +1,199 @@ +# 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]), + pixel_height=_required_int( + scenario[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT + ), + pixel_width=_required_int( + scenario[FIELD_PIXEL_WIDTH], name=FIELD_PIXEL_WIDTH + ), + ), + 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/provider.py b/apps/t2v_app/t2v_app/provider.py deleted file mode 100644 index 692ff9157..000000000 --- a/apps/t2v_app/t2v_app/provider.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Text-to-video application definition for the FlashDreams app host.""" - -from __future__ import annotations - -import argparse -from collections.abc import Mapping, Sequence -from pathlib import Path -from typing import Any - -from flashdreams_app import ( - AppConfig, - AppRequest, - AppSpec, - Mp4RunSpec, - PipelineAppSpec, - WebRTCRunSpec, -) - -from flashdreams.core.pipeline_presets import load_pipeline_provider -from flashdreams.infra.decoder import StreamingVideoDecoder -from flashdreams.infra.pipeline import StreamInferencePipelineConfig -from flashdreams.runtime import InferenceInput - -from .presets import PipelinePreset, RuntimePresetOptions, load_preset_catalog - -FIELD_PROMPT = "prompt" -FIELD_TOTAL_BLOCKS = "total_blocks" -FIELD_PIXEL_HEIGHT = "pixel_height" -FIELD_PIXEL_WIDTH = "pixel_width" -FIELD_FPS = "fps" - - -def parse_options( - parser: argparse.ArgumentParser, - argv: Sequence[str], -) -> Mapping[str, Any]: - """Parse T2V options with the selected mode's host parser. - - Args: - parser: Parser preconfigured with host-owned presentation options. - argv: Arguments remaining after the provider and mode. - - Returns: - Parsed presentation and T2V options keyed by destination. - """ - 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) - return vars(parser.parse_args(argv)) - - -def create_app_spec(request: AppRequest) -> AppSpec: - """Describe a T2V pipeline application without constructing its runtime. - - Args: - request: Parsed host and T2V command-line values. - - Returns: - Pipeline selection, initial conditioning, and presentation data. - """ - options = request.options - preset_id, preset = _resolve_preset(options) - scenario = _scenario(options, preset.runtime) - pipeline_config = _create_pipeline_config(preset_id, preset) - return AppSpec( - config=_build_app_config(scenario, preset.runtime), - pipeline=PipelineAppSpec( - pipeline_config=pipeline_config, - initialize_cache=_initialize_cache, - ), - run=_build_run_spec(request.mode, options, scenario, preset.runtime), - ) - - -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 _build_app_config( - scenario: Mapping[str, object], - runtime_options: RuntimePresetOptions, -) -> AppConfig: - """Build presentation configuration from the resolved T2V scenario.""" - return AppConfig( - model_id="t2v-app", - fps=_required_int(scenario[FIELD_FPS], name=FIELD_FPS), - output_layout=runtime_options.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 - ), - ) - - -def _build_run_spec( - mode: str, - options: Mapping[str, object], - scenario: Mapping[str, object], - defaults: RuntimePresetOptions, -) -> Mp4RunSpec | WebRTCRunSpec: - """Build the selected presentation mode's session data.""" - initial_input = InferenceInput( - global_conditioning={ - FIELD_PROMPT: scenario[FIELD_PROMPT], - FIELD_PIXEL_HEIGHT: scenario[FIELD_PIXEL_HEIGHT], - FIELD_PIXEL_WIDTH: scenario[FIELD_PIXEL_WIDTH], - } - ) - if mode == "webrtc": - return WebRTCRunSpec(initial_input=initial_input) - if mode == "mp4": - total_steps = options.get(FIELD_TOTAL_BLOCKS) - if total_steps is None: - total_steps = defaults.total_blocks - if total_steps is None: - raise ValueError( - "MP4 mode requires --total-blocks or runtime.total_blocks in " - "the selected preset." - ) - return Mp4RunSpec( - initial_input=initial_input, - total_steps=_required_int(total_steps, name=FIELD_TOTAL_BLOCKS), - ) - raise ValueError(f"Unsupported presentation mode: {mode!r}.") - - -def _initialize_cache(pipeline: Any, inputs: InferenceInput) -> object: - """Bind T2V prompt and dimensions to a new pipeline cache.""" - scenario = _scenario_from_inputs(inputs) - decoder = pipeline.decoder - if not isinstance(decoder, StreamingVideoDecoder): - raise TypeError("T2V pipelines require a StreamingVideoDecoder.") - ratio = decoder.spatial_compression_ratio - pixel_height = _required_int(scenario[FIELD_PIXEL_HEIGHT], name=FIELD_PIXEL_HEIGHT) - pixel_width = _required_int(scenario[FIELD_PIXEL_WIDTH], name=FIELD_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})." - ) - return pipeline.initialize_cache( - text=[str(scenario[FIELD_PROMPT])], - image=None, - height=pixel_height // ratio, - width=pixel_width // ratio, - ) - - -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 _scenario_from_inputs(inputs: InferenceInput) -> Mapping[str, object]: - source = inputs.global_conditioning - required = ( - FIELD_PROMPT, - FIELD_PIXEL_HEIGHT, - FIELD_PIXEL_WIDTH, - ) - missing = tuple(name for name in required if name not in source) - if missing: - raise ValueError(f"Missing T2V global conditioning fields: {missing}.") - scenario = dict(source) - scenario[FIELD_PROMPT] = _resolve_prompt(scenario[FIELD_PROMPT]) - for name in ( - FIELD_PIXEL_HEIGHT, - FIELD_PIXEL_WIDTH, - ): - if _required_int(scenario[name], name=name) <= 0: - raise ValueError(f"{name} must be > 0.") - return scenario - - -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_app_spec", "parse_options"] diff --git a/apps/t2v_app/t2v_app/runtime.py b/apps/t2v_app/t2v_app/runtime.py new file mode 100644 index 000000000..cf7122c1c --- /dev/null +++ b/apps/t2v_app/t2v_app/runtime.py @@ -0,0 +1,93 @@ +# 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 + +from typing import Any + +import torch + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineConfig, +) +from flashdreams.runtime import InferenceInput +from flashdreams_runner import AppConfig, IOHandler, Runtime, Session + +from .session import T2VSession, T2VSessionDefaults + + +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 + + @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 + + 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, + ) + + def destroy(self) -> None: + """Release pipeline weights and accelerator allocator state.""" + pipeline = self._pipeline + self._pipeline = None + self._io_handler = None + if pipeline is None: + return + close = getattr(pipeline, "close", None) + if callable(close): + close() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +__all__ = ["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..47b56c8a7 --- /dev/null +++ b/apps/t2v_app/t2v_app/session.py @@ -0,0 +1,189 @@ +# 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 Mapping +from dataclasses import dataclass +from typing import Any, cast + +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, StepResult +from flashdreams_runner import Session + +FIELD_PROMPT = "prompt" +FIELD_PIXEL_HEIGHT = "pixel_height" +FIELD_PIXEL_WIDTH = "pixel_width" + + +@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.""" + + 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.""" + + +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, + ) -> 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[FIELD_PIXEL_HEIGHT] + pixel_width = scenario[FIELD_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._prompt = scenario[FIELD_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 + + @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 generate(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__}." + ) + return StepResult.from_video_chunk( + step_index=index, + video_chunk=video.detach(), + layout=self._output_layout, + metrics=metrics, + ) + + def destroy(self) -> None: + """Release this session's autoregressive cache.""" + if self._destroyed: + return + self._destroyed = True + cache = self._cache + self._cache = None + close = getattr(cache, "close", None) + if callable(close): + close() + + +def _session_scenario( + defaults: T2VSessionDefaults, + initial_input: InferenceInput, +) -> dict[str, Any]: + values = { + FIELD_PROMPT: defaults.prompt, + FIELD_PIXEL_HEIGHT: defaults.pixel_height, + FIELD_PIXEL_WIDTH: defaults.pixel_width, + } + 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 + for name in (FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH): + value = values[name] + 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 values + + +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__ = ["T2VSession", "T2VSessionDefaults"] diff --git a/apps/t2v_app/tests/test_application.py b/apps/t2v_app/tests/test_application.py new file mode 100644 index 000000000..2bdeb5c0f --- /dev/null +++ b/apps/t2v_app/tests/test_application.py @@ -0,0 +1,253 @@ +# 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) + session = runtime.create_session(InferenceInput()) + assert isinstance(session, T2VSession) + result = session.generate(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 index 28e17600a..3cc1111f1 100644 --- a/apps/t2v_app/tests/test_pipeline_provider.py +++ b/apps/t2v_app/tests/test_pipeline_provider.py @@ -20,7 +20,7 @@ from pathlib import Path import pytest -from t2v_app import provider as app_provider +from t2v_app import application from t2v_app.presets import load_preset_catalog from flashdreams.core.checkpoint.remap import unwrap_generator_state_dict @@ -34,7 +34,7 @@ def test_packaged_yaml_constructs_default_pipeline_config() -> None: catalog = load_preset_catalog() preset_id, preset = catalog.resolve(None) - config = app_provider._create_pipeline_config(preset_id, preset) + config = application._create_pipeline_config(preset_id, preset) assert preset_id == "causal-forcing-wan2.1-t2v-1.3b-chunkwise" assert config.name == preset_id diff --git a/apps/t2v_app/tests/test_provider.py b/apps/t2v_app/tests/test_provider.py deleted file mode 100644 index fcc0ba8eb..000000000 --- a/apps/t2v_app/tests/test_provider.py +++ /dev/null @@ -1,141 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""CPU tests for the T2V application provider boundary.""" - -from __future__ import annotations - -import argparse -from typing import Any, cast - -import pytest -import t2v_app -from flashdreams_app import ( - AppProvider, - AppRequest, - AppSpec, - Mp4RunSpec, - PipelineAppSpec, - WebRTCRunSpec, -) -from t2v_app import provider -from t2v_app.presets import ( - PipelinePreset, - PresetCatalog, - RuntimePresetOptions, -) - -from flashdreams.infra.pipeline import StreamInferencePipelineConfig - -pytestmark = pytest.mark.ci_cpu - - -def test_provider_module_conforms_to_host_contract() -> None: - assert isinstance(t2v_app, AppProvider) - - -def test_t2v_provider_parses_model_options() -> None: - parser = argparse.ArgumentParser() - options = provider.parse_options(parser, ["--prompt", "A waterfall"]) - assert options["prompt"] == "A waterfall" - assert options["preset_id"] is None - assert options["preset_config"] is None - assert "backend" not in options - assert "compile" not in options - - -def test_create_app_spec_returns_data_without_constructing_pipeline( - monkeypatch: pytest.MonkeyPatch, -) -> None: - pipeline_constructed = False - - class Pipeline: - def __init__(self, _: object) -> None: - nonlocal pipeline_constructed - pipeline_constructed = True - - preset_id = "test-t2v" - defaults = RuntimePresetOptions( - prompt="default prompt", - total_blocks=2, - pixel_height=64, - pixel_width=96, - fps=12, - output_layout="tchw", - ) - preset = PipelinePreset( - provider="tests:provider", - runtime=defaults, - pipeline={}, - ) - catalog = PresetCatalog( - default_preset_id=preset_id, - presets={preset_id: preset}, - ) - pipeline_config = StreamInferencePipelineConfig( - _target=cast(Any, Pipeline), - name=preset_id, - diffusion_model=cast(Any, None), - ) - - class Provider: - def create_pipeline_config( - self, *, preset_id: str, options: object - ) -> StreamInferencePipelineConfig: - assert preset_id == "test-t2v" - assert options == {} - return pipeline_config - - monkeypatch.setattr(provider, "load_preset_catalog", lambda _: catalog) - monkeypatch.setattr(provider, "load_pipeline_provider", lambda _: Provider()) - - created = provider.create_app_spec( - AppRequest( - mode="mp4", - options={ - "preset_config": None, - "preset_id": None, - "prompt": "A waterfall", - "total_blocks": None, - "pixel_height": None, - "pixel_width": None, - "fps": None, - }, - ) - ) - - assert isinstance(created, AppSpec) - assert isinstance(created.pipeline, PipelineAppSpec) - assert created.pipeline.pipeline_config is pipeline_config - assert created.config.video_width == 96 - assert created.config.fps == 12 - assert isinstance(created.run, Mp4RunSpec) - assert created.run.initial_input.global_conditioning["prompt"] == "A waterfall" - assert created.run.total_steps == 2 - assert not pipeline_constructed - - -def test_webrtc_run_spec_does_not_require_total_steps() -> None: - defaults = RuntimePresetOptions( - prompt="default prompt", - pixel_height=64, - pixel_width=96, - fps=12, - output_layout="tchw", - ) - scenario = { - provider.FIELD_PROMPT: "A waterfall", - provider.FIELD_PIXEL_HEIGHT: 64, - provider.FIELD_PIXEL_WIDTH: 96, - provider.FIELD_FPS: 12, - } - - run_spec = provider._build_run_spec( - "webrtc", - {"total_blocks": None}, - scenario, - defaults, - ) - - assert isinstance(run_spec, WebRTCRunSpec) - assert run_spec.initial_input.global_conditioning["prompt"] == "A waterfall" diff --git a/flashdreams_runner/README.md b/flashdreams_runner/README.md new file mode 100644 index 000000000..195f3ad36 --- /dev/null +++ b/flashdreams_runner/README.md @@ -0,0 +1,152 @@ +# 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 --prompt "A waterfall" +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 |<------| generate(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: + +- `generate(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. + +## 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 one application session per admitted client. | +| `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 generate(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..fb0f20687 --- /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.generate(inputs) + if not isinstance(result, StepResult): + raise TypeError( + "Session.generate() 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..f8b88cf58 --- /dev/null +++ b/flashdreams_runner/contracts.py @@ -0,0 +1,260 @@ +# 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 the runner-facing contract stays create/generate/destroy. + 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 generate(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: + """Describe the next iteration to shared FlashDreams drivers.""" + 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 step(self, inputs: InferenceInput) -> StepResult: + """Generate through the shared inference-session API.""" + result = self.generate(inputs) + if not isinstance(result, StepResult): + raise TypeError( + "Session.generate() must return StepResult, got " + f"{type(result).__name__}." + ) + return result + + 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/apps/flashdreams_app/pyproject.toml b/flashdreams_runner/pyproject.toml similarity index 61% rename from apps/flashdreams_app/pyproject.toml rename to flashdreams_runner/pyproject.toml index 511814b94..16690f24a 100644 --- a/apps/flashdreams_app/pyproject.toml +++ b/flashdreams_runner/pyproject.toml @@ -6,17 +6,18 @@ requires = ["setuptools>=69", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "flashdreams-app" +name = "flashdreams-runner" version = "0.1.0" -description = "Generic host for FlashDreams application runtimes" +description = "Generic shell for FlashDreams applications and I/O modes" requires-python = ">=3.10" dependencies = ["flashdreams[serving]"] [project.scripts] -flashdreams-app = "flashdreams_app.cli:main" +flashdreams-runner = "flashdreams_runner.cli:main" [tool.uv.sources] flashdreams = { workspace = true } -[tool.setuptools.packages.find] -where = ["."] +[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..7a2cb5570 --- /dev/null +++ b/flashdreams_runner/tests/test_cli.py @@ -0,0 +1,306 @@ +# 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 generate(self, inputs: InferenceInput) -> StepResult: + assert not inputs.global_conditioning + calls.append("session.generate") + 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.generate", + "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 generate(self, inputs: InferenceInput) -> StepResult: + calls.append("generate") + 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 + assert session.next_step_request().step_index == 3 + assert session.step(InferenceInput()).step_index == 3 + session.close() + runtime.close() + assert calls == [ + "create_session", + "generate", + "session.destroy", + "runtime.destroy", + ] diff --git a/apps/flashdreams_app/tests/test_webrtc.py b/flashdreams_runner/tests/test_webrtc.py similarity index 54% rename from apps/flashdreams_app/tests/test_webrtc.py rename to flashdreams_runner/tests/test_webrtc.py index c1e4b66c3..6aba328f3 100644 --- a/apps/flashdreams_app/tests/test_webrtc.py +++ b/flashdreams_runner/tests/test_webrtc.py @@ -1,40 +1,54 @@ # 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 import pytest -from flashdreams_app import AppConfig, webrtc -from flashdreams.runtime import InferenceInput, StepRequest, StepResult +from flashdreams.runtime import InferenceInput, StepResult +from flashdreams_runner import AppConfig, IOHandler, Runtime, Session, webrtc pytestmark = pytest.mark.ci_cpu -class _Session: - def next_step_request(self) -> StepRequest | None: - return None - - def step(self, inputs: InferenceInput) -> StepResult: - raise AssertionError("Server construction must not step a session.") +class _Session(Session): + @property + def step_index(self) -> int: + return 0 - def reset(self, inputs: InferenceInput | None = None) -> None: + def generate(self, inputs: InferenceInput) -> StepResult: del inputs + raise AssertionError("Server construction must not generate a chunk.") - def close(self) -> None: + def destroy(self) -> None: pass -class _Runtime: - def start_session(self, inputs: InferenceInput) -> _Session: - del inputs +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 close(self) -> None: + def destroy(self) -> None: pass -def test_host_constructs_webrtc_presentation( +def test_webrtc_mode_constructs_shared_presentation( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} @@ -45,21 +59,10 @@ def fake_serve(**kwargs: object) -> str: monkeypatch.setattr(webrtc, "serve_webrtc_demo", fake_serve) runtime = _Runtime() - initial_input = InferenceInput() result = webrtc.serve_webrtc( runtime=runtime, - config=AppConfig( - model_id="fake-app", - fps=16, - output_layout="tchw", - video_width=96, - video_height=64, - ), - initial_input=initial_input, - options=webrtc.WebRTCOptions( - host="127.0.0.1", - port=8080, - ), + host="127.0.0.1", + port=8080, device="cpu", world_rank=0, ) @@ -69,14 +72,10 @@ def fake_serve(**kwargs: object) -> str: assert captured["world_rank"] == 0 output = captured["output"] assert isinstance(output, webrtc.WebRTCOutputSpec) + assert output.video_width == 96 assert output.warmup_chunks == 0 - assert output.client_liveness_timeout_s == 30.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.runtime_config is output - assert session_manager._shared_scenario is not None - assert session_manager._shared_scenario.initial_inputs is initial_input - assert callable(session_manager._shared_model_input_provider_factory) diff --git a/apps/flashdreams_app/flashdreams_app/webrtc.py b/flashdreams_runner/webrtc.py similarity index 57% rename from apps/flashdreams_app/flashdreams_app/webrtc.py rename to flashdreams_runner/webrtc.py index 1690e2b53..d402c7712 100644 --- a/apps/flashdreams_app/flashdreams_app/webrtc.py +++ b/flashdreams_runner/webrtc.py @@ -1,13 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Host-owned WebRTC presentation for application runtimes.""" +"""WebRTC I/O mode for application runtimes.""" from __future__ import annotations from dataclasses import dataclass +from typing import cast -from flashdreams.runtime import InferenceConfig, InferenceInput, InferenceRuntime +from flashdreams.runtime import ( + InferenceConfig, + InferenceInput, + InferenceRuntime, + OutputArtifact, +) from flashdreams.runtime.demo import ( DemoSpec, PreparedScenario, @@ -21,15 +27,14 @@ 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 AppConfig +from .contracts import DriveSession, Runtime -# TODO: Move this contract into shared FlashDreams and expand it when the -# serving API needs caller-configurable transport and encoder tuning. @dataclass(frozen=True, slots=True) -class WebRTCOptions: - """Minimal WebRTC bind settings for the application prototype.""" +class WebRTCMode: + """Serve application sessions over WebRTC.""" host: str """Server bind address.""" @@ -37,57 +42,83 @@ class WebRTCOptions: 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.""" + + 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, + ) + return () + class _InputProvider: + """Provide empty transport inputs to application-owned sessions.""" + 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 self._initial_input + """Let the application runtime apply its session defaults.""" + return InferenceInput() 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: + """Discard reset inputs because resets create fresh app sessions.""" del inputs def close(self) -> None: - pass + """Release provider resources.""" def serve_webrtc( *, - runtime: InferenceRuntime, - config: AppConfig, - initial_input: InferenceInput, - options: WebRTCOptions, + runtime: Runtime, + host: str, + port: int, device: str, world_rank: int, ) -> object: - """Serve an application runtime through the shared WebRTC stack. + """Serve an initialized application runtime through shared WebRTC. Args: runtime: Initialized application runtime. - config: Model identity and video presentation configuration. - initial_input: Global conditioning used to start live sessions. - options: WebRTC bind settings. + host: Server bind address. + port: Server bind port. device: Device used by the runtime. world_rank: Distributed rank responsible for presentation. Returns: Serving backend result. """ + config = runtime.config output = WebRTCOutputSpec( - host=options.host, - port=options.port, + host=host, + port=port, fps=int(config.fps), video_width=config.video_width, video_height=config.video_height, @@ -99,21 +130,22 @@ def serve_webrtc( output=output, config=InferenceConfig(model_id=config.model_id, device=device), ) - scenario = PreparedScenario(initial_inputs=initial_input) + scenario = PreparedScenario(initial_inputs=InferenceInput()) def create_model_input_provider( spec: DemoSpec, scenario: PreparedScenario, ) -> _InputProvider: del spec, scenario - return _InputProvider(initial_input) + return _InputProvider() + inference_runtime = cast(InferenceRuntime, runtime) manager = BaseWebRTCSessionManager( - runtime=runtime, - runtime_config=output, + runtime=inference_runtime, + runtime_config=cast(WebRTCRuntimeConfig, cast(object, output)), fps=int(config.fps), identity=config.model_id, - shared_host=RuntimeHost(runtime), + shared_host=RuntimeHost(inference_runtime), shared_spec=spec, shared_scenario=scenario, shared_model_input_provider_factory=create_model_input_provider, @@ -126,3 +158,6 @@ def create_model_input_provider( app_resources=WebRTCAppResources(preload_name=config.model_id), world_rank=world_rank, ) + + +__all__ = ["WebRTCMode", "serve_webrtc"] diff --git a/pyproject.toml b/pyproject.toml index 803133510..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,7 +46,6 @@ no-build-isolation-package = ["transformer-engine-torch"] extraPaths = [ "flashdreams", "apps", - "apps/flashdreams_app", "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", @@ -72,7 +72,6 @@ python-version = "3.10" extra-paths = [ "flashdreams", "apps", - "apps/flashdreams_app", "apps/t2v_app", "integrations/omnidreams", "integrations/omnidreams/ludus-renderer", diff --git a/uv.lock b/uv.lock index e275dc9f7..af4874da6 100644 --- a/uv.lock +++ b/uv.lock @@ -20,7 +20,6 @@ conflicts = [[ [manifest] members = [ "flashdreams", - "flashdreams-app", "flashdreams-causal-forcing", "flashdreams-cosmos-predict2", "flashdreams-fastvideo-causal-wan22", @@ -28,6 +27,7 @@ members = [ "flashdreams-hy-worldplay", "flashdreams-lingbot", "flashdreams-omnidreams", + "flashdreams-runner", "flashdreams-sana-wm", "flashdreams-self-forcing", "flashdreams-t2v-demo", @@ -1109,17 +1109,6 @@ cuda13 = [ { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, ] -[[package]] -name = "flashdreams-app" -version = "0.1.0" -source = { editable = "apps/flashdreams_app" } -dependencies = [ - { name = "flashdreams", extra = ["serving"] }, -] - -[package.metadata] -requires-dist = [{ name = "flashdreams", extras = ["serving"], editable = "flashdreams" }] - [[package]] name = "flashdreams-causal-forcing" version = "0.1.0" @@ -1369,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" @@ -4521,13 +4521,13 @@ version = "0.1.0" source = { editable = "apps/t2v_app" } dependencies = [ { name = "flashdreams" }, - { name = "flashdreams-app" }, + { name = "flashdreams-runner" }, ] [package.metadata] requires-dist = [ { name = "flashdreams", editable = "flashdreams" }, - { name = "flashdreams-app", editable = "apps/flashdreams_app" }, + { name = "flashdreams-runner", editable = "flashdreams_runner" }, ] [[package]] From 452786630a28034227fb253bc63bd667dc2908fa Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 17:52:24 +0000 Subject: [PATCH 09/10] Add customizable T2V WebRTC experience Signed-off-by: Gangzheng Tong --- apps/t2v_app/README.md | 9 +- apps/t2v_app/pyproject.toml | 2 +- apps/t2v_app/t2v_app/application.py | 2 + apps/t2v_app/t2v_app/runtime.py | 85 ++++++- apps/t2v_app/t2v_app/session.py | 134 ++++++++-- apps/t2v_app/t2v_app/web/adapter.css | 10 + apps/t2v_app/t2v_app/web/adapter.js | 31 +++ apps/t2v_app/t2v_app/webrtc.py | 240 ++++++++++++++++++ apps/t2v_app/tests/test_application.py | 1 + apps/t2v_app/tests/test_webrtc.py | 234 +++++++++++++++++ .../flashdreams/serving/webrtc/manager.py | 7 +- flashdreams_runner/README.md | 10 +- flashdreams_runner/contracts.py | 4 +- flashdreams_runner/tests/test_cli.py | 4 +- flashdreams_runner/tests/test_webrtc.py | 50 ++++ flashdreams_runner/webrtc.py | 132 ++++++++-- 16 files changed, 898 insertions(+), 57 deletions(-) create mode 100644 apps/t2v_app/t2v_app/web/adapter.css create mode 100644 apps/t2v_app/t2v_app/web/adapter.js create mode 100644 apps/t2v_app/t2v_app/webrtc.py create mode 100644 apps/t2v_app/tests/test_webrtc.py diff --git a/apps/t2v_app/README.md b/apps/t2v_app/README.md index 7145e356a..df99daa0d 100644 --- a/apps/t2v_app/README.md +++ b/apps/t2v_app/README.md @@ -22,8 +22,7 @@ uv run flashdreams-runner t2v-app mp4 \ --output outputs/waterfall.mp4 uv run flashdreams-runner t2v-app webrtc \ - --preset-id self-forcing-wan2.1-t2v-1.3b \ - --prompt "A neon-lit city at night" + --preset-id self-forcing-wan2.1-t2v-1.3b uv run flashdreams-runner t2v-app none \ --steps 2 \ @@ -40,8 +39,10 @@ construct a `StreamInferencePipelineConfig`. The packaged catalog is 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. WebRTC does not impose a finite -step count. +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: diff --git a/apps/t2v_app/pyproject.toml b/apps/t2v_app/pyproject.toml index 7694afd3e..531e67e0a 100644 --- a/apps/t2v_app/pyproject.toml +++ b/apps/t2v_app/pyproject.toml @@ -26,4 +26,4 @@ venv = ".venv" where = ["."] [tool.setuptools.package-data] -t2v_app = ["pipeline_presets.yaml"] +t2v_app = ["pipeline_presets.yaml", "web/*.css", "web/*.js"] diff --git a/apps/t2v_app/t2v_app/application.py b/apps/t2v_app/t2v_app/application.py index a747efea7..9923e8c31 100644 --- a/apps/t2v_app/t2v_app/application.py +++ b/apps/t2v_app/t2v_app/application.py @@ -68,12 +68,14 @@ def create_runtime(arguments: ApplicationArguments) -> Runtime: 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", diff --git a/apps/t2v_app/t2v_app/runtime.py b/apps/t2v_app/t2v_app/runtime.py index cf7122c1c..53edaa73d 100644 --- a/apps/t2v_app/t2v_app/runtime.py +++ b/apps/t2v_app/t2v_app/runtime.py @@ -17,7 +17,10 @@ from __future__ import annotations -from typing import Any +import math +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast import torch @@ -27,8 +30,20 @@ ) from flashdreams.runtime import InferenceInput from flashdreams_runner import AppConfig, IOHandler, Runtime, Session +from flashdreams_runner.webrtc import WebRTCMode -from .session import T2VSession, T2VSessionDefaults +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): @@ -46,6 +61,8 @@ def __init__( 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: @@ -64,6 +81,11 @@ def initialize(self, *, device: str, io_handler: IOHandler) -> None: ) 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.""" @@ -74,13 +96,70 @@ def create_session(self, initial_input: InferenceInput | None = None) -> Session 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) @@ -90,4 +169,4 @@ def destroy(self) -> None: torch.cuda.empty_cache() -__all__ = ["T2VRuntime"] +__all__ = ["T2VArtifact", "T2VRuntime"] diff --git a/apps/t2v_app/t2v_app/session.py b/apps/t2v_app/t2v_app/session.py index 47b56c8a7..aded8a4ea 100644 --- a/apps/t2v_app/t2v_app/session.py +++ b/apps/t2v_app/t2v_app/session.py @@ -17,21 +17,46 @@ from __future__ import annotations -from collections.abc import Mapping +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, StepResult +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) @@ -41,12 +66,18 @@ class T2VSessionDefaults: 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.""" @@ -58,14 +89,16 @@ def __init__( 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[FIELD_PIXEL_HEIGHT] - pixel_width = scenario[FIELD_PIXEL_WIDTH] + 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 " @@ -73,7 +106,8 @@ def __init__( ) self._pipeline = pipeline - self._prompt = scenario[FIELD_PROMPT] + self._scenario = scenario + self._prompt = scenario.prompt self._pixel_height = pixel_height self._pixel_width = pixel_width self._output_layout: VideoTensorLayout = output_layout @@ -87,6 +121,19 @@ def __init__( 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: @@ -98,6 +145,20 @@ 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 generate(self, inputs: InferenceInput) -> StepResult: """Generate and finalize one autoregressive video block.""" if self._destroyed or self._cache is None: @@ -126,12 +187,16 @@ def generate(self, inputs: InferenceInput) -> StepResult: "T2V pipeline generate() must return torch.Tensor, got " f"{type(video).__name__}." ) - return StepResult.from_video_chunk( + 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.""" @@ -140,32 +205,65 @@ def destroy(self) -> None: self._destroyed = True cache = self._cache self._cache = None - close = getattr(cache, "close", None) - if callable(close): - close() + 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, -) -> dict[str, Any]: +) -> 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 - for name in (FIELD_PIXEL_HEIGHT, FIELD_PIXEL_WIDTH): - value = values[name] - 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 values + 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]: @@ -186,4 +284,4 @@ def _metrics(value: object) -> Mapping[str, float | int]: return metrics -__all__ = ["T2VSession", "T2VSessionDefaults"] +__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 index 2bdeb5c0f..5296a601d 100644 --- a/apps/t2v_app/tests/test_application.py +++ b/apps/t2v_app/tests/test_application.py @@ -219,6 +219,7 @@ def run(self, runtime: Runtime, drive_session: DriveSession) -> tuple[()]: 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.generate(InferenceInput()) diff --git a/apps/t2v_app/tests/test_webrtc.py b/apps/t2v_app/tests/test_webrtc.py new file mode 100644 index 000000000..6d5c41d92 --- /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.generate(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/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index db22b87a7..b0785be07 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/flashdreams/flashdreams/serving/webrtc/manager.py @@ -686,6 +686,7 @@ def __init__( 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") @@ -702,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() @@ -1353,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 diff --git a/flashdreams_runner/README.md b/flashdreams_runner/README.md index 195f3ad36..1b4ec172f 100644 --- a/flashdreams_runner/README.md +++ b/flashdreams_runner/README.md @@ -8,7 +8,7 @@ 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 --prompt "A waterfall" +uv run flashdreams-runner t2v-app webrtc uv run flashdreams-runner t2v-app none --steps 4 --prompt "A waterfall" ``` @@ -64,6 +64,12 @@ 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 @@ -73,7 +79,7 @@ implement application generation logic. |---|---| | `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 one application session per admitted client. | +| `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`. diff --git a/flashdreams_runner/contracts.py b/flashdreams_runner/contracts.py index f8b88cf58..6c2f52f83 100644 --- a/flashdreams_runner/contracts.py +++ b/flashdreams_runner/contracts.py @@ -179,8 +179,8 @@ 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: - """Describe the next iteration to shared FlashDreams drivers.""" + 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 diff --git a/flashdreams_runner/tests/test_cli.py b/flashdreams_runner/tests/test_cli.py index 7a2cb5570..b46926489 100644 --- a/flashdreams_runner/tests/test_cli.py +++ b/flashdreams_runner/tests/test_cli.py @@ -294,7 +294,9 @@ def destroy(self) -> None: runtime = FakeRuntime() assert runtime.start_session(InferenceInput()) is session - assert session.next_step_request().step_index == 3 + 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() diff --git a/flashdreams_runner/tests/test_webrtc.py b/flashdreams_runner/tests/test_webrtc.py index 6aba328f3..f4c57ae21 100644 --- a/flashdreams_runner/tests/test_webrtc.py +++ b/flashdreams_runner/tests/test_webrtc.py @@ -5,9 +5,12 @@ 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 @@ -79,3 +82,50 @@ def fake_serve(**kwargs: object) -> str: 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 index d402c7712..661be7178 100644 --- a/flashdreams_runner/webrtc.py +++ b/flashdreams_runner/webrtc.py @@ -5,8 +5,9 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import cast +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any, Protocol, cast from flashdreams.runtime import ( InferenceConfig, @@ -16,6 +17,7 @@ ) from flashdreams.runtime.demo import ( DemoSpec, + ModelInputProvider, PreparedScenario, PreparedStep, ProviderCapabilities, @@ -32,7 +34,7 @@ from .contracts import DriveSession, Runtime -@dataclass(frozen=True, slots=True) +@dataclass(slots=True) class WebRTCMode: """Serve application sessions over WebRTC.""" @@ -51,6 +53,18 @@ class WebRTCMode: 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, @@ -64,21 +78,25 @@ def run( port=self.port, device=self.device, world_rank=self.world_rank, + customization=self._customization, ) return () class _InputProvider: - """Provide empty transport inputs to application-owned sessions.""" + """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: - """Let the application runtime apply its session defaults.""" - return InferenceInput() + """Return the application-provided input for a new session.""" + return self._initial_input def prepare_step( self, *, request: StepRequirements, user_window: UserInputWindow @@ -88,13 +106,53 @@ def prepare_step( return PreparedStep(inference_input=InferenceInput()) def reset(self, inputs: InferenceInput | None = None) -> None: - """Discard reset inputs because resets create fresh app sessions.""" - del inputs + """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, @@ -102,6 +160,7 @@ def serve_webrtc( port: int, device: str, world_rank: int, + customization: WebRTCCustomization | None = None, ) -> object: """Serve an initialized application runtime through shared WebRTC. @@ -111,6 +170,7 @@ def serve_webrtc( 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. @@ -130,34 +190,58 @@ def serve_webrtc( output=output, config=InferenceConfig(model_id=config.model_id, device=device), ) - scenario = PreparedScenario(initial_inputs=InferenceInput()) + 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, scenario - return _InputProvider() + del spec + return _InputProvider(scenario.initial_inputs) inference_runtime = cast(InferenceRuntime, runtime) - manager = 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, - ) + 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=WebRTCAppResources(preload_name=config.model_id), + app_resources=app_resources, world_rank=world_rank, ) -__all__ = ["WebRTCMode", "serve_webrtc"] +__all__ = [ + "ModelInputProviderFactory", + "WebRTCCustomization", + "WebRTCMode", + "serve_webrtc", +] From e96e5e1754aaeeddd4a9418b96097f5ab37f6255 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 13 Aug 2026 18:40:33 +0000 Subject: [PATCH 10/10] Standardize application sessions on step Use the shared inference session spelling directly to remove the duplicate generate adapter and keep runner and WebRTC execution aligned. Signed-off-by: Gangzheng Tong --- PR_DESCRIPTION.md | 6 +- apps/t2v_app/ARCHITECTURE.md | 133 ++++++++++++++++++++++++ apps/t2v_app/t2v_app/session.py | 2 +- apps/t2v_app/tests/test_application.py | 2 +- apps/t2v_app/tests/test_webrtc.py | 2 +- flashdreams_runner/README.md | 6 +- flashdreams_runner/cli.py | 4 +- flashdreams_runner/contracts.py | 14 +-- flashdreams_runner/tests/test_cli.py | 12 +-- flashdreams_runner/tests/test_webrtc.py | 4 +- 10 files changed, 154 insertions(+), 31 deletions(-) create mode 100644 apps/t2v_app/ARCHITECTURE.md diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md index 51d4273d3..d621a0d33 100644 --- a/PR_DESCRIPTION.md +++ b/PR_DESCRIPTION.md @@ -24,7 +24,7 @@ uv run flashdreams-runner t2v-app {mp4 | replay | webrtc | none} | | input / output | Session | | main loop: | ---------------------------> | prompt + cache | | read input | <--------------------------- | generate/finalize | -| Session.generate | StepResult +--------------------+ +| Session.step | StepResult +--------------------+ | present output | +------------+-------------+ | @@ -91,7 +91,7 @@ When `--preset-id` is omitted, `t2v-app` uses the catalog's `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.generate()` and `Session.destroy()` for per-user prompt, +- 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 @@ -114,7 +114,7 @@ When `--preset-id` is omitted, `t2v-app` uses the catalog's 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.generate()`. + inside `T2VSession.step()`. - Keep prompts, dimensions, caches, and step indexes isolated per session. ## Shared FlashDreams changes 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/t2v_app/session.py b/apps/t2v_app/t2v_app/session.py index aded8a4ea..901856b09 100644 --- a/apps/t2v_app/t2v_app/session.py +++ b/apps/t2v_app/t2v_app/session.py @@ -159,7 +159,7 @@ def next_step_request(self) -> StepRequest | None: }, ) - def generate(self, inputs: InferenceInput) -> StepResult: + 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.") diff --git a/apps/t2v_app/tests/test_application.py b/apps/t2v_app/tests/test_application.py index 5296a601d..cfef9418f 100644 --- a/apps/t2v_app/tests/test_application.py +++ b/apps/t2v_app/tests/test_application.py @@ -222,7 +222,7 @@ def run(self, runtime: Runtime, drive_session: DriveSession) -> tuple[()]: assert runtime.peek_steady_output_num_frames() == 3 session = runtime.create_session(InferenceInput()) assert isinstance(session, T2VSession) - result = session.generate(InferenceInput()) + result = session.step(InferenceInput()) assert result.step_index == 0 assert result.frame_count == 3 assert result.metrics["step_ms"] == 1.0 diff --git a/apps/t2v_app/tests/test_webrtc.py b/apps/t2v_app/tests/test_webrtc.py index 6d5c41d92..45d9f73fb 100644 --- a/apps/t2v_app/tests/test_webrtc.py +++ b/apps/t2v_app/tests/test_webrtc.py @@ -224,7 +224,7 @@ def close(self) -> tuple[OutputArtifact, ...]: ) assert session.next_step_request() is not None - session.generate(InferenceInput()) + session.step(InferenceInput()) assert session.next_step_request() is None session.destroy() diff --git a/flashdreams_runner/README.md b/flashdreams_runner/README.md index 1b4ec172f..35f3d45ef 100644 --- a/flashdreams_runner/README.md +++ b/flashdreams_runner/README.md @@ -23,7 +23,7 @@ application module flashdreams-runner I/O mode | model weights | | loop: | +-------+-------+ | global state | | input = mode.read() | | | | | output = Session. |<--------------+ -| Session |<------| generate(input) | +| Session |<------| step(input) | | prompt/cache | | mode.write(output) |-------------->+ | game state | | destroy Session/Runtime | +------------------+ +--------------------------+ @@ -56,7 +56,7 @@ package and `flashdreams-runner`. `Session` owns the application loop implementation and all per-user state, such as prompts, K/V caches, world state, and step counters: -- `generate(inputs)` performs exactly one application iteration and returns a +- `step(inputs)` performs exactly one application iteration and returns a `StepResult`. - `destroy()` releases session-local resources. @@ -110,7 +110,7 @@ class MySession(Session): def step_index(self) -> int: return self._step_index - def generate(self, inputs: InferenceInput) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: video = self.model.generate(self.cache, inputs.step) result = StepResult.from_video_chunk( step_index=self._step_index, diff --git a/flashdreams_runner/cli.py b/flashdreams_runner/cli.py index fb0f20687..626c80d93 100644 --- a/flashdreams_runner/cli.py +++ b/flashdreams_runner/cli.py @@ -203,10 +203,10 @@ def close_output() -> None: resources.callback(session.destroy) while (inputs := input_handler.read()) is not None: - result = session.generate(inputs) + result = session.step(inputs) if not isinstance(result, StepResult): raise TypeError( - "Session.generate() must return StepResult, got " + "Session.step() must return StepResult, got " f"{type(result).__name__}." ) output_handler.write(result) diff --git a/flashdreams_runner/contracts.py b/flashdreams_runner/contracts.py index 6c2f52f83..fe6102e72 100644 --- a/flashdreams_runner/contracts.py +++ b/flashdreams_runner/contracts.py @@ -143,7 +143,7 @@ def destroy(self) -> None: """Release model weights and process-wide resources.""" # These aliases let shared FlashDreams serving code consume the application - # ABI directly while the runner-facing contract stays create/generate/destroy. + # 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) @@ -171,7 +171,7 @@ def steady_output_frame_count(self) -> int | None: return None @abstractmethod - def generate(self, inputs: InferenceInput) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: """Run one application main-loop iteration.""" @abstractmethod @@ -186,16 +186,6 @@ def next_step_request(self) -> StepRequest | None: metadata["steady_output_frame_count"] = self.steady_output_frame_count return StepRequest(step_index=self.step_index, metadata=metadata) - def step(self, inputs: InferenceInput) -> StepResult: - """Generate through the shared inference-session API.""" - result = self.generate(inputs) - if not isinstance(result, StepResult): - raise TypeError( - "Session.generate() must return StepResult, got " - f"{type(result).__name__}." - ) - return result - def reset(self, inputs: InferenceInput | None = None) -> None: """Reject reset when an application requires a fresh session.""" del inputs diff --git a/flashdreams_runner/tests/test_cli.py b/flashdreams_runner/tests/test_cli.py index b46926489..a8980f3ae 100644 --- a/flashdreams_runner/tests/test_cli.py +++ b/flashdreams_runner/tests/test_cli.py @@ -73,9 +73,9 @@ def __init__(self) -> None: def step_index(self) -> int: return self._step_index - def generate(self, inputs: InferenceInput) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: assert not inputs.global_conditioning - calls.append("session.generate") + calls.append("session.step") result = _result(self._step_index) self._step_index += 1 return result @@ -182,7 +182,7 @@ def run( "input.initial_input", "runtime.create_session", "input.read", - "session.generate", + "session.step", "output.write", "input.read", "output.close", @@ -266,8 +266,8 @@ class FakeSession(Session): def step_index(self) -> int: return 3 - def generate(self, inputs: InferenceInput) -> StepResult: - calls.append("generate") + def step(self, inputs: InferenceInput) -> StepResult: + calls.append("step") return _result(3) def destroy(self) -> None: @@ -302,7 +302,7 @@ def destroy(self) -> None: runtime.close() assert calls == [ "create_session", - "generate", + "step", "session.destroy", "runtime.destroy", ] diff --git a/flashdreams_runner/tests/test_webrtc.py b/flashdreams_runner/tests/test_webrtc.py index f4c57ae21..289dbd095 100644 --- a/flashdreams_runner/tests/test_webrtc.py +++ b/flashdreams_runner/tests/test_webrtc.py @@ -21,9 +21,9 @@ class _Session(Session): def step_index(self) -> int: return 0 - def generate(self, inputs: InferenceInput) -> StepResult: + def step(self, inputs: InferenceInput) -> StepResult: del inputs - raise AssertionError("Server construction must not generate a chunk.") + raise AssertionError("Server construction must not run a step.") def destroy(self) -> None: pass