diff --git a/fastvideo/api/omni.py b/fastvideo/api/omni.py new file mode 100644 index 000000000..507faf775 --- /dev/null +++ b/fastvideo/api/omni.py @@ -0,0 +1,624 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Typed omni request plane (design.md §6.1). + +This module introduces the request-plane vocabulary the next-generation +runtime is built on: + +- :class:`OmniRequest` — a typed multimodal request whose ``task`` is + *declared, never inferred*, and whose inputs are typed + :class:`ModalPart` s instead of per-model fields on a god-object. +- :class:`OmniOutput` — a typed multimodal output whose modalities are + named :class:`Artifact` slots carrying provenance, replacing the + ``extra["audio"]`` escape hatch (design.md P3). +- :data:`OmniEvent` — one streaming-event union (progress / chunk / + final), the single channel that ``LoopStage.step`` emits through. + +It is deliberately additive and engine-agnostic: nothing here imports +``torch`` or the pipeline machinery, and adapters bridge to/from today's +:class:`~fastvideo.api.schema.GenerationRequest` / +:class:`~fastvideo.api.results.GenerationResult` so the new types are +usable through the existing ``VideoGenerator`` before the engine itself +is rebuilt (plan.md M1). Later milestones evolve ``GenerationRequest`` +into ``OmniRequest`` in place and drop the adapters. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, ClassVar +from uuid import uuid4 + +from fastvideo.api.results import ( + GenerationResult, + VideoEvent, + VideoFinalEvent, + VideoPartialEvent, + VideoProgressEvent, +) +from fastvideo.api.schema import ( + GenerationRequest, + InputConfig, + OutputConfig, + RequestRuntimeConfig, + SamplingConfig, +) + + +class Modality(str, Enum): + """A media modality. ``str`` mixin keeps it JSON/serialization friendly.""" + + TEXT = "text" + IMAGE = "image" + VIDEO = "video" + AUDIO = "audio" + ACTION = "action" + LATENT = "latent" + + +class TaskType(str, Enum): + """The declared task (design.md §6.1 / kills P7). + + The pipeline graph branches on ``request.task``. Heuristics may only + *suggest* a default at the API boundary (see :func:`infer_task`); they + never decide control flow inside the runtime. + """ + + T2V = "t2v" # text -> video + I2V = "i2v" # image (+text) -> video + TI2V = "ti2v" # text + init image -> video + V2V = "v2v" # video -> video (edit / restyle) + V2W = "v2w" # video -> world (continue a world-model rollout) + T2I = "t2i" # text -> image + I2I = "i2i" # image -> image (edit) + T2A = "t2a" # text -> audio + T2VS = "t2vs" # text -> video + sound (joint A/V) + A2W = "a2w" # action -> world (interactive world model) + REASON = "reason" # text -> text (AR reasoner / prompt upsampling) + + +# --------------------------------------------------------------------------- +# Inputs: typed modality parts (replaces InputConfig's per-model fields, P3). +# --------------------------------------------------------------------------- + + +@dataclass +class ModalPart: + """Base for a typed input part. + + ``modality`` is intrinsic to the concrete subclass (a ``ClassVar``, not + an instance field). ``role`` disambiguates several parts of one + modality — e.g. ``"prompt"`` vs ``"negative"`` text, ``"init"`` vs + ``"conditioning"`` image — and is keyword-only so subclass payloads stay + positional. + """ + + modality: ClassVar[Modality] + role: str | None = field(default=None, kw_only=True) + + +@dataclass +class TextPart(ModalPart): + modality: ClassVar[Modality] = Modality.TEXT + text: str = "" + + +@dataclass +class ImagePart(ModalPart): + modality: ClassVar[Modality] = Modality.IMAGE + image: Any | None = None # PIL.Image / ndarray / tensor + path: str | None = None + + +@dataclass +class VideoPart(ModalPart): + modality: ClassVar[Modality] = Modality.VIDEO + video: Any | None = None + path: str | None = None + fps: float | None = None + + +@dataclass +class AudioPart(ModalPart): + modality: ClassVar[Modality] = Modality.AUDIO + audio: Any | None = None + path: str | None = None + sample_rate: int | None = None + + +@dataclass +class ActionPart(ModalPart): + modality: ClassVar[Modality] = Modality.ACTION + action: Any | None = None # mouse / keyboard / camera tensors + kind: str | None = None # "mouse" | "keyboard" | "camera" | ... + + +@dataclass +class LatentPart(ModalPart): + modality: ClassVar[Modality] = Modality.LATENT + latents: Any | None = None + of_modality: Modality = Modality.VIDEO # modality these latents decode to + + +# --------------------------------------------------------------------------- +# Per-call parameters: AR sampling vs diffusion, separated (design.md §6.1). +# --------------------------------------------------------------------------- + + +@dataclass +class SamplingParams: + """AR decode knobs. + + Distinct from the legacy diffusion ``fastvideo.api.SamplingParam``: these + drive ``ARDecodeLoop`` (Cosmos3 reasoner, omni thinkers/talkers, codec + decode), not denoising. + """ + + max_tokens: int = 512 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int | None = None + stop: list[str] = field(default_factory=list) + seed: int | None = None + + +@dataclass +class DiffusionParams: + """Denoise knobs. ``guidance_per_modality`` carries per-modality CFG + scales for joint A/V denoise (LTX-2, Cosmos3 t2vs); ``guidance_scale`` is + the scalar default.""" + + steps: int = 50 + guidance_scale: float = 1.0 + guidance_per_modality: dict[Modality, float] = field(default_factory=dict) + sigmas: list[float] | None = None + flow_shift: float | None = None + height: int | None = None + width: int | None = None + num_frames: int | None = None + fps: int | None = None + seed: int | None = None + + +# --------------------------------------------------------------------------- +# Outputs spec: requested modalities + streaming + capture flags. +# --------------------------------------------------------------------------- + + +@dataclass +class StreamSpec: + """Per-modality streaming policy. ``chunk_ms`` applies to audio chunks, + ``per_chunk`` to chunked-causal video; ``enabled`` alone covers token + text.""" + + enabled: bool = True + chunk_ms: int | None = None + per_chunk: bool = False + + +@dataclass +class OutputSpec: + """Requested output modalities, streaming, and capture flags.""" + + modalities: list[Modality] = field(default_factory=lambda: [Modality.VIDEO]) + stream: dict[Modality, StreamSpec] = field(default_factory=dict) + return_latents: bool = False + return_trajectory: bool = False + + @property + def streaming(self) -> bool: + """True if any modality is requested as a stream.""" + return any(spec.enabled for spec in self.stream.values()) + + +@dataclass +class NodeOverrides: + """Per-graph-node parameter overrides (design.md §6.1). + + For a multi-loop graph, ``node_params["refine"].steps`` overrides the + refine loop's step count without leaking ``refine_*`` onto the universal + schema. Validation against each node's declared schema arrives with + ``PipelineSpec`` (a later milestone); for now this is a typed bag with + ``get`` / item / attribute access. + """ + + params: dict[str, Any] = field(default_factory=dict) + + def get(self, key: str, default: Any = None) -> Any: + return self.params.get(key, default) + + def __getitem__(self, key: str) -> Any: + return self.params[key] + + def __getattr__(self, key: str) -> Any: + # Only invoked when normal attribute lookup fails, so ``params`` + # itself resolves through ``__dict__`` and never recurses. + try: + return self.__dict__["params"][key] + except KeyError as exc: + raise AttributeError(key) from exc + + +@dataclass +class OmniRequest: + """A typed multimodal request (design.md §6.1).""" + + task: TaskType + inputs: list[ModalPart] = field(default_factory=list) + sampling: SamplingParams = field(default_factory=SamplingParams) + diffusion: DiffusionParams = field(default_factory=DiffusionParams) + outputs: OutputSpec = field(default_factory=OutputSpec) + node_params: dict[str, NodeOverrides] = field(default_factory=dict) + priority: int = 0 + request_id: str = field(default_factory=lambda: uuid4().hex) + + # -- accessors ---------------------------------------------------------- + + def parts(self, modality: Modality) -> list[ModalPart]: + return [p for p in self.inputs if p.modality is modality] + + @property + def prompt(self) -> str | None: + for part in self.inputs: + if isinstance(part, TextPart) and part.role in (None, "prompt"): + return part.text + return None + + @property + def negative_prompt(self) -> str | None: + for part in self.inputs: + if isinstance(part, TextPart) and part.role in ("negative", "negative_prompt"): + return part.text + return None + + # -- constructors / adapters ------------------------------------------- + + @classmethod + def from_prompt( + cls, + prompt: str | None, + task: TaskType, + *, + negative_prompt: str | None = None, + **kwargs: Any, + ) -> OmniRequest: + """Build a request from a bare prompt — what ``generate_video`` calls + internally so the offline shim constructs an ``OmniRequest`` (G5).""" + inputs: list[ModalPart] = [] + if prompt is not None: + inputs.append(TextPart(prompt)) + if negative_prompt is not None: + inputs.append(TextPart(negative_prompt, role="negative")) + return cls(task=task, inputs=inputs, **kwargs) + + @classmethod + def from_generation_request( + cls, + request: GenerationRequest, + task: TaskType | None = None, + ) -> OmniRequest: + """Lift a legacy typed :class:`GenerationRequest` into an + ``OmniRequest``; ``task`` defaults to the boundary heuristic.""" + resolved = task if task is not None else infer_task(request) + inputs: list[ModalPart] = [] + prompt = request.prompt + if isinstance(prompt, list): + prompt = prompt[0] if prompt else None + if prompt is not None: + inputs.append(TextPart(prompt)) + if request.negative_prompt is not None: + inputs.append(TextPart(request.negative_prompt, role="negative")) + + inp = request.inputs + image_path = inp.image_path if isinstance(inp.image_path, str) else None + if image_path is not None or inp.pil_image is not None: + inputs.append(ImagePart(image=inp.pil_image, path=image_path)) + video_path = inp.video_path if isinstance(inp.video_path, str) else None + if video_path is not None: + inputs.append(VideoPart(path=video_path)) + if inp.mouse_cond is not None or inp.keyboard_cond is not None: + inputs.append(ActionPart(action=inp.mouse_cond, kind="mouse")) + + sampling = request.sampling + diffusion = DiffusionParams( + steps=sampling.num_inference_steps, + guidance_scale=sampling.guidance_scale, + sigmas=sampling.sigmas, + height=sampling.height, + width=sampling.width, + num_frames=sampling.num_frames, + fps=sampling.fps, + seed=sampling.seed, + ) + outputs = OutputSpec( + return_trajectory=request.runtime.return_trajectory_latents, + return_latents=request.runtime.return_trajectory_decoded, + ) + node_params = { + node: NodeOverrides(params=dict(overrides)) + for node, overrides in request.stage_overrides.items() if isinstance(overrides, dict) + } + return cls( + task=resolved, + inputs=inputs, + diffusion=diffusion, + outputs=outputs, + node_params=node_params, + ) + + def to_generation_request(self) -> GenerationRequest: + """Lower to a legacy :class:`GenerationRequest` so today's + ``VideoGenerator`` can execute an ``OmniRequest`` unchanged (M1).""" + diff = self.diffusion + sampling = SamplingConfig() + sampling.num_inference_steps = diff.steps + sampling.guidance_scale = diff.guidance_scale + sampling.sigmas = diff.sigmas + if diff.seed is not None: + sampling.seed = diff.seed + if diff.num_frames is not None: + sampling.num_frames = diff.num_frames + if diff.height is not None: + sampling.height = diff.height + if diff.width is not None: + sampling.width = diff.width + if diff.fps is not None: + sampling.fps = diff.fps + + image_path: str | None = None + pil_image: Any | None = None + video_path: str | None = None + for part in self.inputs: + if isinstance(part, ImagePart): + image_path = image_path or part.path + pil_image = pil_image if pil_image is not None else part.image + elif isinstance(part, VideoPart): + video_path = video_path or part.path + inputs = InputConfig(image_path=image_path, pil_image=pil_image, video_path=video_path) + + runtime = RequestRuntimeConfig( + return_trajectory_latents=self.outputs.return_trajectory, + return_trajectory_decoded=self.outputs.return_latents, + ) + output = OutputConfig(return_frames=Modality.VIDEO in self.outputs.modalities) + stage_overrides = {node: dict(ov.params) for node, ov in self.node_params.items()} + return GenerationRequest( + prompt=self.prompt, + negative_prompt=self.negative_prompt, + inputs=inputs, + sampling=sampling, + runtime=runtime, + output=output, + stage_overrides=stage_overrides, + ) + + +# --------------------------------------------------------------------------- +# Outputs: named artifacts with provenance (kills extra["audio"], P3). +# --------------------------------------------------------------------------- + + +@dataclass +class Artifact: + """Base output artifact. ``source_node`` records which graph node + produced it (provenance, design.md §6.1).""" + + modality: ClassVar[Modality] + source_node: str | None = field(default=None, kw_only=True) + + +@dataclass +class VideoArtifact(Artifact): + modality: ClassVar[Modality] = Modality.VIDEO + frames: Any | None = None # numpy (N, H, W, 3) uint8 + tensor: Any | None = None # raw sample tensor + path: str | None = None + fps: float | None = None + + +@dataclass +class AudioArtifact(Artifact): + modality: ClassVar[Modality] = Modality.AUDIO + audio: Any | None = None + sample_rate: int | None = None + + +@dataclass +class TextArtifact(Artifact): + modality: ClassVar[Modality] = Modality.TEXT + text: str = "" + token_ids: list[int] | None = None + + +@dataclass +class TensorArtifact(Artifact): + """Action tensors and other raw tensor outputs.""" + + modality: ClassVar[Modality] = Modality.ACTION + tensor: Any | None = None + + +@dataclass +class LatentArtifact(Artifact): + modality: ClassVar[Modality] = Modality.LATENT + latents: Any | None = None + timesteps: Any | None = None + of_modality: Modality = Modality.VIDEO + + +@dataclass +class RequestMetrics: + generation_time: float | None = None + peak_memory_mb: float | None = None + extra: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class OmniOutput: + """Typed multimodal output (design.md §6.1).""" + + request_id: str + artifacts: dict[str, Artifact] = field(default_factory=dict) + metrics: RequestMetrics = field(default_factory=RequestMetrics) + + def get(self, name: str) -> Artifact | None: + return self.artifacts.get(name) + + @property + def video(self) -> VideoArtifact | None: + art = self.artifacts.get("video") + return art if isinstance(art, VideoArtifact) else None + + @property + def audio(self) -> AudioArtifact | None: + art = self.artifacts.get("audio") + return art if isinstance(art, AudioArtifact) else None + + @classmethod + def from_generation_result( + cls, + result: GenerationResult, + request_id: str = "", + ) -> OmniOutput: + """Map a legacy :class:`GenerationResult` into named artifacts. + + Audio becomes a first-class :class:`AudioArtifact` carrying its + sample rate, instead of riding in ``extra["audio"]`` (P3). + """ + artifacts: dict[str, Artifact] = {} + if result.frames is not None or result.samples is not None or result.video_path is not None: + artifacts["video"] = VideoArtifact( + frames=result.frames, + tensor=result.samples, + path=result.video_path, + source_node="decode", + ) + if result.audio is not None: + artifacts["audio"] = AudioArtifact( + audio=result.audio, + sample_rate=result.audio_sample_rate, + source_node="audio_decode", + ) + if result.trajectory is not None or result.trajectory_decoded is not None: + artifacts["latents"] = LatentArtifact( + latents=result.trajectory, + timesteps=result.trajectory_timesteps, + source_node="denoise", + ) + return cls( + request_id=request_id, + artifacts=artifacts, + metrics=RequestMetrics( + generation_time=result.generation_time, + peak_memory_mb=result.peak_memory_mb, + ), + ) + + +# --------------------------------------------------------------------------- +# Streaming: one event union (evolves api/results.py's Video*Event). +# --------------------------------------------------------------------------- + + +@dataclass +class OmniProgressEvent: + """Per-step progress telemetry.""" + + step: int + total_steps: int + node: str = "denoise" + + +@dataclass +class OmniChunkEvent: + """A streamed artifact chunk — the universal ``StepResult.emit`` channel + (design.md §6.2.2): text tokens, audio chunks, or decoded frame chunks. + + ``pts`` optionally carries a presentation timestamp for raw-frame + streaming over WebRTC (design.md §9.1).""" + + modality: Modality + index: int + payload: Any = None + pts: float | None = None + node: str | None = None + + +@dataclass +class OmniFinalEvent: + """Terminal event carrying the full :class:`OmniOutput`.""" + + output: OmniOutput + + +OmniEvent = OmniProgressEvent | OmniChunkEvent | OmniFinalEvent +"""Union of every event the engine streams; consumers match by ``isinstance``.""" + + +def infer_task(request: GenerationRequest) -> TaskType: + """Best-effort boundary heuristic for a legacy request (design.md §6.1). + + A *suggestion* only — the runtime always branches on the explicit + ``OmniRequest.task``, never on this. Single-frame requests are images; + a video input implies edit; an image input implies image-to-video. + """ + inp = request.inputs + has_image = bool(inp.image_path or inp.pil_image) + has_video = bool(inp.video_path or inp.stage1_video) + if request.sampling.num_frames == 1: + return TaskType.I2I if has_image else TaskType.T2I + if has_video: + return TaskType.V2V + if has_image: + return TaskType.I2V + return TaskType.T2V + + +def omni_event_from_video_event(event: VideoEvent, request_id: str = "") -> OmniEvent: + """Adapt a legacy :data:`~fastvideo.api.results.VideoEvent` to an + :data:`OmniEvent` so the streaming surface can migrate incrementally.""" + if isinstance(event, VideoProgressEvent): + return OmniProgressEvent(step=event.step, total_steps=event.total_steps, node=event.stage) + if isinstance(event, VideoPartialEvent): + return OmniChunkEvent(modality=Modality.VIDEO, index=event.index, payload=event.frames, node="decode") + if isinstance(event, VideoFinalEvent): + if event.result is not None: + output = OmniOutput.from_generation_result(event.result, request_id) + else: + output = OmniOutput(request_id=request_id) + if event.frames is not None: + output.artifacts["video"] = VideoArtifact(frames=event.frames, source_node="decode") + return OmniFinalEvent(output=output) + raise TypeError(f"unknown VideoEvent type: {type(event).__name__}") + + +__all__ = [ + "ActionPart", + "Artifact", + "AudioArtifact", + "AudioPart", + "DiffusionParams", + "ImagePart", + "LatentArtifact", + "LatentPart", + "Modality", + "ModalPart", + "NodeOverrides", + "OmniChunkEvent", + "OmniEvent", + "OmniFinalEvent", + "OmniOutput", + "OmniProgressEvent", + "OmniRequest", + "OutputSpec", + "RequestMetrics", + "SamplingParams", + "StreamSpec", + "TaskType", + "TensorArtifact", + "TextArtifact", + "TextPart", + "VideoArtifact", + "VideoPart", + "infer_task", + "omni_event_from_video_event", +] diff --git a/fastvideo/tests/api/test_omni.py b/fastvideo/tests/api/test_omni.py new file mode 100644 index 000000000..e6f40aeea --- /dev/null +++ b/fastvideo/tests/api/test_omni.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the typed omni request plane (fastvideo/api/omni.py, design.md §6.1).""" +from __future__ import annotations + +import pytest + +from fastvideo.api.omni import ( + AudioArtifact, + DiffusionParams, + ImagePart, + Modality, + NodeOverrides, + OmniChunkEvent, + OmniFinalEvent, + OmniOutput, + OmniProgressEvent, + OmniRequest, + OutputSpec, + StreamSpec, + TaskType, + TextPart, + VideoArtifact, + infer_task, + omni_event_from_video_event, +) +from fastvideo.api.results import ( + GenerationResult, + VideoFinalEvent, + VideoPartialEvent, + VideoProgressEvent, +) +from fastvideo.api.schema import GenerationRequest, InputConfig, SamplingConfig + + +def test_enums_are_strings(): + assert TaskType.T2V == "t2v" + assert Modality.AUDIO == "audio" + # str mixin keeps them usable anywhere a string is expected + assert f"{TaskType.REASON}" == "TaskType.REASON" or TaskType.REASON.value == "reason" + assert TaskType("i2v") is TaskType.I2V + + +def test_request_id_autogenerated_and_unique(): + a = OmniRequest(task=TaskType.T2V) + b = OmniRequest(task=TaskType.T2V) + assert a.request_id and b.request_id + assert a.request_id != b.request_id + + +def test_prompt_and_negative_accessors(): + req = OmniRequest.from_prompt("a cat", TaskType.T2V, negative_prompt="blurry") + assert req.prompt == "a cat" + assert req.negative_prompt == "blurry" + assert req.parts(Modality.TEXT) + assert req.parts(Modality.IMAGE) == [] + + +def test_positional_text_part_is_text_not_role(): + # role is keyword-only, so the positional arg is the payload + part = TextPart("hello") + assert part.text == "hello" + assert part.role is None + + +def test_round_trip_through_generation_request(): + req = OmniRequest( + task=TaskType.T2V, + inputs=[TextPart("a fox"), TextPart("ugly", role="negative")], + diffusion=DiffusionParams(steps=30, guidance_scale=6.0, height=480, width=832, num_frames=49), + ) + legacy = req.to_generation_request() + assert isinstance(legacy, GenerationRequest) + assert legacy.prompt == "a fox" + assert legacy.negative_prompt == "ugly" + assert legacy.sampling.num_inference_steps == 30 + assert legacy.sampling.guidance_scale == 6.0 + assert legacy.sampling.height == 480 + assert legacy.sampling.num_frames == 49 + + back = OmniRequest.from_generation_request(legacy, task=TaskType.T2V) + assert back.prompt == "a fox" + assert back.negative_prompt == "ugly" + assert back.diffusion.steps == 30 + assert back.diffusion.guidance_scale == 6.0 + assert back.diffusion.height == 480 + assert back.diffusion.num_frames == 49 + + +def test_image_input_lowers_and_lifts(): + req = OmniRequest(task=TaskType.I2V, inputs=[TextPart("dance"), ImagePart(path="/tmp/x.png")]) + legacy = req.to_generation_request() + assert legacy.inputs.image_path == "/tmp/x.png" + back = OmniRequest.from_generation_request(legacy) + assert back.task == TaskType.I2V # inferred from the image input + assert back.parts(Modality.IMAGE)[0].path == "/tmp/x.png" # type: ignore[attr-defined] + + +@pytest.mark.parametrize( + ("inputs", "sampling", "expected"), + [ + (InputConfig(), SamplingConfig(num_frames=49), TaskType.T2V), + (InputConfig(image_path="/a.png"), SamplingConfig(num_frames=49), TaskType.I2V), + (InputConfig(video_path="/a.mp4"), SamplingConfig(num_frames=49), TaskType.V2V), + (InputConfig(), SamplingConfig(num_frames=1), TaskType.T2I), + (InputConfig(image_path="/a.png"), SamplingConfig(num_frames=1), TaskType.I2I), + ], +) +def test_infer_task_heuristic(inputs, sampling, expected): + req = GenerationRequest(prompt="x", inputs=inputs, sampling=sampling) + assert infer_task(req) == expected + + +def test_named_artifacts_replace_extra_audio(): + result = GenerationResult( + prompt="song", + frames="FRAMES", + audio="WAV", + audio_sample_rate=44100, + generation_time=1.5, + peak_memory_mb=1234.0, + ) + out = OmniOutput.from_generation_result(result, request_id="r1") + assert out.request_id == "r1" + assert isinstance(out.video, VideoArtifact) + assert out.video.frames == "FRAMES" + # audio is a first-class artifact carrying its sample rate, not extra["audio"] + assert isinstance(out.audio, AudioArtifact) + assert out.audio.sample_rate == 44100 + assert out.audio.source_node == "audio_decode" + assert out.metrics.generation_time == 1.5 + assert out.metrics.peak_memory_mb == 1234.0 + + +def test_omni_event_from_video_event(): + prog = omni_event_from_video_event(VideoProgressEvent(step=3, total_steps=50, stage="refine")) + assert isinstance(prog, OmniProgressEvent) + assert prog.step == 3 and prog.node == "refine" + + chunk = omni_event_from_video_event(VideoPartialEvent(frames="CHUNK", index=2)) + assert isinstance(chunk, OmniChunkEvent) + assert chunk.modality == Modality.VIDEO and chunk.index == 2 and chunk.payload == "CHUNK" + + final = omni_event_from_video_event( + VideoFinalEvent(result=GenerationResult(frames="F", audio="A", audio_sample_rate=24000)), + request_id="rid", + ) + assert isinstance(final, OmniFinalEvent) + assert final.output.request_id == "rid" + assert final.output.audio.sample_rate == 24000 + + +def test_node_overrides_access(): + req = OmniRequest(task=TaskType.T2V, node_params={"refine": NodeOverrides(params={"steps": 8})}) + assert req.node_params["refine"].steps == 8 # attribute access + assert req.node_params["refine"]["steps"] == 8 # item access + assert req.node_params["refine"].get("missing", 0) == 0 + with pytest.raises(AttributeError): + _ = req.node_params["refine"].nonexistent + # node_params survive the round-trip into stage_overrides + legacy = req.to_generation_request() + assert legacy.stage_overrides == {"refine": {"steps": 8}} + + +def test_output_spec_streaming_flag(): + spec = OutputSpec(modalities=[Modality.VIDEO, Modality.AUDIO]) + assert spec.streaming is False + spec.stream[Modality.AUDIO] = StreamSpec(enabled=True, chunk_ms=200) + assert spec.streaming is True