diff --git a/flashdreams/flashdreams/infra/runner_io.py b/flashdreams/flashdreams/infra/runner_io.py new file mode 100644 index 000000000..465f98e88 --- /dev/null +++ b/flashdreams/flashdreams/infra/runner_io.py @@ -0,0 +1,353 @@ +# 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-facing image, video, prompt, and stats I/O helpers.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any, Literal, TypeAlias + +import numpy as np +import torch + +IMAGE_SUFFIXES = frozenset({".bmp", ".jpeg", ".jpg", ".png", ".webp"}) +"""Image filename suffixes treated as still images by runner helpers.""" + +DEFAULT_RUNNER_INSTALL_HINT = ( + "Install the runner extras: pip install 'flashdreams[runners]'." +) +"""Default install hint for optional runner I/O dependencies.""" + +ResizeInterpolation: TypeAlias = Literal[ + "default", "nearest", "linear", "area", "cubic", "lanczos4" +] +"""OpenCV resize interpolation names accepted by runner image/video helpers.""" + +VideoTensorLayout: TypeAlias = Literal["thwc", "tchw", "bcthw"] +"""Tensor layouts accepted by ``write_video_tensor``.""" + +InputAssetValidator: TypeAlias = Callable[[Path], object] +"""Validator signature for runner input assets downloaded to local cache.""" + + +def ensure_output_dir(output_dir: Path) -> Path: + """Create and return a runner output directory.""" + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir + + +def runner_artifact_path(output_dir: Path, runner_name: str, suffix: str) -> Path: + """Return the conventional artifact path for a runner output.""" + return output_dir / f"{runner_name}.{suffix.lstrip('.')}" + + +def runner_stats_path(output_dir: Path, runner_name: str) -> Path: + """Return the conventional per-runner stats JSON path.""" + return output_dir / f"stats_{runner_name}.json" + + +def write_runner_stats( + output_dir: Path, + runner_name: str, + stats: Any, + *, + indent: int = 2, +) -> Path: + """Write runner stats JSON and return the output path.""" + path = runner_stats_path(output_dir, runner_name) + path.write_text(json.dumps(stats, indent=indent)) + return path + + +def resolve_prompt_value(value: str | Path) -> str: + """Resolve an inline prompt or first non-empty line of a prompt file.""" + if isinstance(value, Path): + lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] + if not lines: + raise ValueError(f"prompt file {value} has no non-empty lines") + return lines[0] + if not value: + raise ValueError("--prompt must be a non-empty string or a path to a .txt file") + return value + + +def resolve_input_path( + value: str | Path, + *, + cache_dir: Path, + filename: str | None = None, + validator: InputAssetValidator | None = None, +) -> Path: + """Resolve a local input path or download an HTTP(S) asset into cache.""" + if isinstance(value, Path): + return value + if not value.startswith(("http://", "https://")): + return Path(value) + return _download_to_cache( + value, + cache_dir=cache_dir, + filename=filename, + validator=validator, + ) + + +def _download_to_cache( + url: str, + *, + cache_dir: Path, + filename: str | None, + validator: InputAssetValidator | None, +) -> Path: + """Download an asset through the core downloader without importing it eagerly.""" + from flashdreams.core.io.download import download_to_cache # noqa: PLC0415 + + return download_to_cache( + url, + cache_dir=cache_dir, + filename=filename, + validator=validator, + ) + + +def read_image_rgb( + path: str | Path, + *, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> np.ndarray: + """Read an RGB image as ``[H, W, 3]``.""" + media = _import_mediapy("Loading images", install_hint=install_hint) + return media.read_image(str(path))[..., :3] + + +def read_video_rgb( + path: str | Path, + *, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> np.ndarray: + """Read an RGB video as ``[T, H, W, 3]``.""" + media = _import_mediapy("Loading videos", install_hint=install_hint) + return media.read_video(str(path))[..., :3] + + +def read_video_fps( + path: str | Path, + *, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> float: + """Read a video's frame rate from ``mediapy`` metadata.""" + media = _import_mediapy("Probing video metadata", install_hint=install_hint) + return float(media.VideoMetadata.from_path(str(path)).fps) + + +def read_first_frame_rgb( + path: str | Path, + *, + image_suffixes: set[str] | frozenset[str] = IMAGE_SUFFIXES, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> np.ndarray: + """Read an image or the first frame of a video as ``[H, W, 3]``.""" + path = Path(path) + media = _import_mediapy("Loading first-frame assets", install_hint=install_hint) + if path.suffix.lower() in image_suffixes: + return media.read_image(str(path))[..., :3] + video = media.read_video(str(path)) + if video.shape[0] == 0: + raise ValueError(f"video has no frames: {path}") + return video[0, ..., :3] + + +def resize_rgb_image( + image: np.ndarray, + *, + pixel_height: int, + pixel_width: int, + interpolation: ResizeInterpolation = "default", + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> np.ndarray: + """Resize an RGB image with OpenCV's ``(width, height)`` convention.""" + cv2 = _import_cv2("Resizing images", install_hint=install_hint) + resize_kwargs = _cv2_resize_kwargs(cv2, interpolation) + return cv2.resize(image, (pixel_width, pixel_height), **resize_kwargs) + + +def resize_rgb_video( + video: np.ndarray, + *, + pixel_height: int, + pixel_width: int, + interpolation: ResizeInterpolation = "default", + skip_if_matching: bool = True, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> np.ndarray: + """Resize an RGB video frame-by-frame with OpenCV.""" + if skip_if_matching and video.shape[1:3] == (pixel_height, pixel_width): + return video + return np.stack( + [ + resize_rgb_image( + frame, + pixel_height=pixel_height, + pixel_width=pixel_width, + interpolation=interpolation, + install_hint=install_hint, + ) + for frame in video + ], + axis=0, + ) + + +def rgb_image_to_normalized_tensor( + image: np.ndarray, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Convert an RGB image to ``[1, C, H, W]`` in ``[-1, 1]``.""" + tensor = torch.from_numpy(image).to(device=device, dtype=dtype) / 127.5 - 1.0 + return tensor.permute(2, 0, 1).unsqueeze(0) + + +def rgb_video_to_normalized_tensor( + video: np.ndarray, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Convert an RGB video to ``[T, C, H, W]`` in ``[-1, 1]``.""" + tensor = torch.from_numpy(video).to(device=device, dtype=dtype) / 127.5 - 1.0 + return tensor.permute(0, 3, 1, 2) + + +def load_first_frame_tensor( + path: str | Path, + *, + pixel_height: int, + pixel_width: int, + device: torch.device, + dtype: torch.dtype, + allow_video: bool = False, + interpolation: ResizeInterpolation = "default", + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> torch.Tensor: + """Load, resize, and normalize a first-frame asset.""" + if allow_video: + image = read_first_frame_rgb(path, install_hint=install_hint) + else: + image = read_image_rgb(path, install_hint=install_hint) + image = resize_rgb_image( + image, + pixel_height=pixel_height, + pixel_width=pixel_width, + interpolation=interpolation, + install_hint=install_hint, + ) + return rgb_image_to_normalized_tensor(image, device=device, dtype=dtype) + + +def load_video_tensor( + path: str | Path, + *, + pixel_height: int, + pixel_width: int, + device: torch.device, + dtype: torch.dtype, + interpolation: ResizeInterpolation = "default", + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> torch.Tensor: + """Load, resize, and normalize an RGB video.""" + video = read_video_rgb(path, install_hint=install_hint) + video = resize_rgb_video( + video, + pixel_height=pixel_height, + pixel_width=pixel_width, + interpolation=interpolation, + install_hint=install_hint, + ) + return rgb_video_to_normalized_tensor(video, device=device, dtype=dtype) + + +def video_tensor_to_uint8( + video: torch.Tensor, + *, + layout: VideoTensorLayout, +) -> np.ndarray: + """Convert a ``[-1, 1]`` video tensor to uint8 ``[T, H, W, C]`` frames.""" + if layout == "thwc": + canvas = video + elif layout == "tchw": + canvas = video.permute(0, 2, 3, 1) + elif layout == "bcthw": + if video.shape[0] != 1: + raise ValueError( + "layout='bcthw' expects a single batch element; " + f"got {tuple(video.shape)}" + ) + canvas = video[0].permute(1, 2, 3, 0) + else: + raise ValueError(f"Unsupported video tensor layout: {layout!r}") + + arr = (canvas.detach().cpu().float().numpy() + 1.0) / 2.0 + return (arr * 255).clip(0, 255).astype("uint8") + + +def write_video_tensor( + video: torch.Tensor, + path: str | Path, + *, + fps: int | float, + layout: VideoTensorLayout, + install_hint: str = DEFAULT_RUNNER_INSTALL_HINT, +) -> Path: + """Write a ``[-1, 1]`` video tensor as an MP4 and return the path.""" + media = _import_mediapy("Writing videos", install_hint=install_hint) + path = Path(path) + media.write_video(str(path), video_tensor_to_uint8(video, layout=layout), fps=fps) + return path + + +def _import_mediapy(action: str, *, install_hint: str) -> Any: + """Import ``mediapy`` with a runner-specific dependency hint.""" + try: + import mediapy as media # noqa: PLC0415 + except ImportError as exc: # pragma: no cover - import-time gate + raise ImportError(f"{action} needs mediapy. {install_hint}") from exc + return media + + +def _import_cv2(action: str, *, install_hint: str) -> Any: + """Import ``cv2`` with a runner-specific dependency hint.""" + try: + import cv2 # noqa: PLC0415 + except ImportError as exc: # pragma: no cover - import-time gate + raise ImportError(f"{action} needs opencv. {install_hint}") from exc + return cv2 + + +def _cv2_resize_kwargs(cv2: Any, interpolation: ResizeInterpolation) -> dict[str, int]: + """Return OpenCV resize kwargs for the requested interpolation.""" + if interpolation == "default": + return {} + interpolation_values = { + "nearest": cv2.INTER_NEAREST, + "linear": cv2.INTER_LINEAR, + "area": cv2.INTER_AREA, + "cubic": cv2.INTER_CUBIC, + "lanczos4": cv2.INTER_LANCZOS4, + } + return {"interpolation": interpolation_values[interpolation]} diff --git a/flashdreams/flashdreams/recipes/template/runner.py b/flashdreams/flashdreams/recipes/template/runner.py index 12394a233..f5516e0d5 100644 --- a/flashdreams/flashdreams/recipes/template/runner.py +++ b/flashdreams/flashdreams/recipes/template/runner.py @@ -34,6 +34,7 @@ from flashdreams.infra.pipeline import StreamInferencePipeline from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ensure_output_dir, runner_artifact_path from flashdreams.recipes.template.encoder import TemplateControlEncoder from flashdreams.recipes.template.transformer import ( TemplateTransformer, @@ -135,8 +136,8 @@ def run(self) -> None: if not self.is_rank_zero: return stacked = torch.stack(outputs, dim=0).cpu() - cfg.output_dir.mkdir(parents=True, exist_ok=True) - out_path = cfg.output_dir / f"{cfg.runner_name}.pt" + ensure_output_dir(cfg.output_dir) + out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "pt") torch.save(stacked, out_path) logger.info( diff --git a/flashdreams/tests/test_runner_io.py b/flashdreams/tests/test_runner_io.py new file mode 100644 index 000000000..30a61d4d0 --- /dev/null +++ b/flashdreams/tests/test_runner_io.py @@ -0,0 +1,308 @@ +# 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 runner I/O helpers.""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch + +import flashdreams.infra.runner_io as runner_io +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + read_first_frame_rgb, + read_video_fps, + resolve_input_path, + resolve_prompt_value, + runner_artifact_path, + runner_stats_path, + video_tensor_to_uint8, + write_runner_stats, + write_video_tensor, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_resolve_prompt_value_reads_first_non_empty_line(tmp_path: Path) -> None: + prompt_path = tmp_path / "prompt.txt" + prompt_path.write_text("\n first prompt \nsecond prompt\n") + + assert resolve_prompt_value("inline") == "inline" + assert resolve_prompt_value(prompt_path) == "first prompt" + + +def test_resolve_prompt_value_rejects_empty_values(tmp_path: Path) -> None: + prompt_path = tmp_path / "prompt.txt" + prompt_path.write_text("\n \n") + + with pytest.raises(ValueError, match="has no non-empty lines"): + resolve_prompt_value(prompt_path) + with pytest.raises(ValueError, match="must be a non-empty string"): + resolve_prompt_value("") + + +def test_resolve_input_path_passes_local_paths_through(tmp_path: Path) -> None: + cache_dir = tmp_path / "cache" + + assert resolve_input_path(tmp_path / "frame.png", cache_dir=cache_dir) == ( + tmp_path / "frame.png" + ) + assert resolve_input_path("relative.mp4", cache_dir=cache_dir) == Path( + "relative.mp4" + ) + + +def test_resolve_input_path_downloads_urls( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[tuple[str, Path, str | None, runner_io.InputAssetValidator | None]] = [] + + def validator(path: Path) -> object: + return path + + def fake_download_to_cache( + url: str, + *, + cache_dir: Path, + filename: str | None = None, + validator: runner_io.InputAssetValidator | None = None, + ) -> Path: + calls.append((url, cache_dir, filename, validator)) + return cache_dir / (filename or "asset.bin") + + monkeypatch.setattr(runner_io, "_download_to_cache", fake_download_to_cache) + + resolved = resolve_input_path( + "https://example.test/asset.png", + cache_dir=tmp_path / "cache", + filename="input.png", + validator=validator, + ) + + assert resolved == tmp_path / "cache" / "input.png" + assert calls == [ + ("https://example.test/asset.png", tmp_path / "cache", "input.png", validator) + ] + + +def test_runner_artifact_and_stats_paths(tmp_path: Path) -> None: + output_dir = ensure_output_dir(tmp_path / "nested") + + assert output_dir.is_dir() + assert runner_artifact_path(output_dir, "demo-runner", "mp4") == ( + output_dir / "demo-runner.mp4" + ) + assert runner_artifact_path(output_dir, "demo-runner", ".mp4") == ( + output_dir / "demo-runner.mp4" + ) + assert runner_stats_path(output_dir, "demo-runner") == ( + output_dir / "stats_demo-runner.json" + ) + + +def test_write_runner_stats_matches_existing_json_format(tmp_path: Path) -> None: + stats_path = write_runner_stats( + tmp_path, "demo", [{"autoregressive_index": 0, "total_ms": 12.5}] + ) + + assert stats_path == tmp_path / "stats_demo.json" + assert stats_path.read_text() == ( + '[\n {\n "autoregressive_index": 0,\n "total_ms": 12.5\n }\n]' + ) + + +def test_video_tensor_to_uint8_converts_tchw_layout() -> None: + video = torch.tensor( + [ + [ + [[-1.0, 0.0], [1.0, 2.0]], + [[-2.0, 0.5], [0.0, 1.0]], + [[1.0, -1.0], [0.0, 0.0]], + ], + ], + dtype=torch.float32, + ) + + frames = video_tensor_to_uint8(video, layout="tchw") + + assert frames.dtype == np.uint8 + assert frames.shape == (1, 2, 2, 3) + np.testing.assert_array_equal( + frames, + np.array( + [[[[0, 0, 255], [127, 191, 0]], [[255, 127, 127], [255, 255, 127]]]], + dtype=np.uint8, + ), + ) + + +def test_video_tensor_to_uint8_converts_thwc_layout() -> None: + video = torch.tensor( + [ + [ + [[-1.0, 0.0, 1.0], [0.5, -0.5, 0.0]], + [[1.0, 1.0, -1.0], [2.0, -2.0, 0.0]], + ], + ], + dtype=torch.float32, + ) + + frames = video_tensor_to_uint8(video, layout="thwc") + + assert frames.dtype == np.uint8 + assert frames.shape == (1, 2, 2, 3) + np.testing.assert_array_equal( + frames, + np.array( + [[[[0, 127, 255], [191, 63, 127]], [[255, 255, 0], [255, 0, 127]]]], + dtype=np.uint8, + ), + ) + + +def test_video_tensor_to_uint8_converts_bcthw_layout() -> None: + video = torch.full((1, 3, 2, 4, 5), -1.0, dtype=torch.float32) + + frames = video_tensor_to_uint8(video, layout="bcthw") + + assert frames.shape == (2, 4, 5, 3) + assert frames.dtype == np.uint8 + assert frames.max() == 0 + + +def test_video_tensor_to_uint8_rejects_multi_batch_bcthw() -> None: + video = torch.zeros((2, 3, 1, 4, 5), dtype=torch.float32) + + with pytest.raises(ValueError, match="expects a single batch element"): + video_tensor_to_uint8(video, layout="bcthw") + + +def test_read_first_frame_rgb_rejects_empty_video( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_media = types.ModuleType("mediapy") + + def read_video(path: str) -> np.ndarray: + return np.empty((0, 2, 2, 3), dtype=np.uint8) + + setattr(fake_media, "read_video", read_video) + monkeypatch.setitem(sys.modules, "mediapy", fake_media) + + with pytest.raises(ValueError, match="video has no frames"): + read_first_frame_rgb(Path("empty.mp4")) + + +def test_read_video_fps_uses_lazy_mediapy_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_media = types.ModuleType("mediapy") + calls: list[str] = [] + + class VideoMetadata: + fps = 23.976 + + @classmethod + def from_path(cls, path: str) -> "VideoMetadata": + calls.append(path) + return cls() + + setattr(fake_media, "VideoMetadata", VideoMetadata) + monkeypatch.setitem(sys.modules, "mediapy", fake_media) + + assert read_video_fps(Path("clip.mp4")) == pytest.approx(23.976) + assert calls == ["clip.mp4"] + + +def test_write_video_tensor_lazy_imports_mediapy( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + calls: list[tuple[str, np.ndarray, int]] = [] + fake_media = types.ModuleType("mediapy") + + def write_video(path: str, frames: np.ndarray, *, fps: int) -> None: + calls.append((path, frames, fps)) + + setattr(fake_media, "write_video", write_video) + monkeypatch.setitem(sys.modules, "mediapy", fake_media) + + out_path = tmp_path / "out.mp4" + returned = write_video_tensor( + torch.zeros((1, 3, 2, 2), dtype=torch.float32), + out_path, + fps=16, + layout="tchw", + ) + + assert returned == out_path + assert calls[0][0] == str(out_path) + assert calls[0][1].shape == (1, 2, 2, 3) + assert calls[0][1].dtype == np.uint8 + assert calls[0][2] == 16 + + +def test_load_first_frame_tensor_uses_requested_resize_interpolation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake_media = types.ModuleType("mediapy") + fake_cv2 = types.ModuleType("cv2") + calls: dict[str, Any] = {} + + setattr(fake_cv2, "INTER_CUBIC", 2) + setattr(fake_cv2, "INTER_NEAREST", 0) + setattr(fake_cv2, "INTER_LINEAR", 1) + setattr(fake_cv2, "INTER_AREA", 3) + setattr(fake_cv2, "INTER_LANCZOS4", 4) + + def read_image(path: str) -> np.ndarray: + calls["path"] = path + return np.full((2, 3, 4), 127, dtype=np.uint8) + + def resize(image: np.ndarray, dsize: tuple[int, int], **kwargs: int) -> np.ndarray: + calls["dsize"] = dsize + calls["kwargs"] = kwargs + width, height = dsize + return np.full((height, width, 3), image[0, 0, 0], dtype=image.dtype) + + setattr(fake_media, "read_image", read_image) + setattr(fake_cv2, "resize", resize) + monkeypatch.setitem(sys.modules, "mediapy", fake_media) + monkeypatch.setitem(sys.modules, "cv2", fake_cv2) + + tensor = load_first_frame_tensor( + Path("frame.png"), + pixel_height=4, + pixel_width=5, + device=torch.device("cpu"), + dtype=torch.float32, + interpolation="cubic", + ) + + assert calls == { + "path": "frame.png", + "dsize": (5, 4), + "kwargs": {"interpolation": 2}, + } + assert tensor.shape == (1, 3, 4, 5) + assert tensor.dtype == torch.float32 + assert torch.allclose(tensor, torch.full_like(tensor, 127.0 / 127.5 - 1.0)) diff --git a/integrations/causal_forcing/causal_forcing/runner.py b/integrations/causal_forcing/causal_forcing/runner.py index 5b30e7b93..767b4a452 100644 --- a/integrations/causal_forcing/causal_forcing/runner.py +++ b/integrations/causal_forcing/causal_forcing/runner.py @@ -17,21 +17,26 @@ from __future__ import annotations -import json import os from dataclasses import dataclass, field from pathlib import Path -import cv2 -import mediapy as media import torch -from einops import rearrange from loguru import logger -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + read_image_rgb, + resolve_input_path, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashdreams.recipes.wan import ( WanInferencePipeline, WanInferencePipelineCache, @@ -67,25 +72,6 @@ """User-writable cache for on-the-fly I2V first-frame downloads.""" -def _resolve_image_path(image_path: str | Path) -> Path: - """Return a local ``Path`` for ``image_path``, downloading URLs on the fly. - - ``http(s)://`` strings are atomically fetched into - :data:`IMAGE_CACHE_DIR` and validated as decodable images before - being published; local paths pass through unchanged. - """ - if isinstance(image_path, Path): - return image_path - if not image_path.startswith(("http://", "https://")): - return Path(image_path) - - return download_to_cache( - image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=lambda p: media.read_image(str(p)), - ) - - @dataclass(kw_only=True) class CausalForcingT2VRunnerConfig(RunnerConfig): """Runner config for the Causal-Forcing T2V variants. @@ -151,13 +137,7 @@ def _resolve_prompt(self) -> str: A Path reads its first non-empty line, a str is used as-is. """ - value = self.config.prompt - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value + return resolve_prompt_value(self.config.prompt) def _initialize_cache(self) -> WanInferencePipelineCache: """Initialize the autoregressive cache for T2V.""" @@ -208,13 +188,9 @@ def run(self) -> None: return # Write the video. - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "t c h w -> t h w c") - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=config.fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -223,8 +199,9 @@ def run(self) -> None: # Write the perf stats. if stats_history: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + config.output_dir, config.runner_name, stats_history + ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) @@ -258,14 +235,16 @@ def _initialize_cache(self) -> WanInferencePipelineCache: # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the # pipeline's actual device so non-default ``--device`` selections # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - local_image_path = _resolve_image_path(config.image_path) - arr = media.read_image(str(local_image_path))[..., :3] - arr = cv2.resize(arr, (config.pixel_width, config.pixel_height)) - tensor = ( - torch.from_numpy(arr).to(device=self.pipeline.device, dtype=torch.bfloat16) - / 127.5 - - 1.0 + image = load_first_frame_tensor( + resolve_input_path( + config.image_path, + cache_dir=IMAGE_CACHE_DIR, + validator=read_image_rgb, + ), + pixel_height=config.pixel_height, + pixel_width=config.pixel_width, + device=self.pipeline.device, + dtype=torch.bfloat16, ) - image = rearrange(tensor, "h w c -> 1 c h w") return self.pipeline.initialize_cache(text=[prompt], image=image) diff --git a/integrations/cosmos_predict2/cosmos_predict2/runner.py b/integrations/cosmos_predict2/cosmos_predict2/runner.py index ebf2bcd2d..8aafeb6db 100644 --- a/integrations/cosmos_predict2/cosmos_predict2/runner.py +++ b/integrations/cosmos_predict2/cosmos_predict2/runner.py @@ -17,21 +17,26 @@ from __future__ import annotations -import json import os from dataclasses import dataclass, field from pathlib import Path -import cv2 -import mediapy as media import torch -from einops import rearrange from loguru import logger -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + read_image_rgb, + resolve_input_path, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashdreams.recipes.cosmos.pipeline import ( CosmosInferencePipeline, CosmosInferencePipelineCache, @@ -71,25 +76,6 @@ """User-writable cache for on-the-fly I2V first-frame downloads.""" -def _resolve_image_path(image_path: str | Path) -> Path: - """Return a local ``Path`` for ``image_path``, downloading URLs on the fly. - - ``http(s)://`` strings are atomically fetched into - :data:`IMAGE_CACHE_DIR` and validated as decodable images before - being published; local paths pass through unchanged. - """ - if isinstance(image_path, Path): - return image_path - if not image_path.startswith(("http://", "https://")): - return Path(image_path) - - return download_to_cache( - image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=lambda p: media.read_image(str(p)), - ) - - @dataclass(kw_only=True) class Cosmos2T2VRunnerConfig(RunnerConfig): """Runner config for the Cosmos-Predict2 T2V variant.""" @@ -123,13 +109,7 @@ def _resolve_prompt(self) -> str: A Path reads its first non-empty line, a str is used as-is. """ - value = self.config.prompt - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value + return resolve_prompt_value(self.config.prompt) def _initialize_cache(self) -> CosmosInferencePipelineCache: """Initialize the autoregressive cache for T2V.""" @@ -165,13 +145,9 @@ def run(self) -> None: if generated is None: return - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "t c h w -> t h w c") - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=config.fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -179,9 +155,10 @@ def run(self) -> None: ) if stats is not None: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text( - json.dumps([{"autoregressive_index": 0, **stats}], indent=2) + stats_path = write_runner_stats( + config.output_dir, + config.runner_name, + [{"autoregressive_index": 0, **stats}], ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats " @@ -222,14 +199,16 @@ def _initialize_cache(self) -> CosmosInferencePipelineCache: # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the # pipeline's actual device so non-default ``--device`` selections # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - local_image_path = _resolve_image_path(config.image_path) - arr = media.read_image(str(local_image_path))[..., :3] - arr = cv2.resize(arr, (config.pixel_width, config.pixel_height)) - tensor = ( - torch.from_numpy(arr).to(device=self.pipeline.device, dtype=torch.bfloat16) - / 127.5 - - 1.0 + image = load_first_frame_tensor( + resolve_input_path( + config.image_path, + cache_dir=IMAGE_CACHE_DIR, + validator=read_image_rgb, + ), + pixel_height=config.pixel_height, + pixel_width=config.pixel_width, + device=self.pipeline.device, + dtype=torch.bfloat16, ) - image = rearrange(tensor, "h w c -> 1 c h w") return self.pipeline.initialize_cache(text=[prompt], image=image) diff --git a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py index 151a70275..a66b44b0b 100644 --- a/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py +++ b/integrations/fastvideo_causal_wan22/fastvideo_causal_wan22/runner.py @@ -17,17 +17,21 @@ from __future__ import annotations -import json from dataclasses import dataclass, field from pathlib import Path -import mediapy as media -from einops import rearrange from loguru import logger from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashdreams.recipes.wan import ( WanInferencePipeline, WanInferencePipelineCache, @@ -93,13 +97,7 @@ def _resolve_prompt(self) -> str: A Path reads its first non-empty line, a str is used as-is. """ - value = self.config.prompt - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value + return resolve_prompt_value(self.config.prompt) def _initialize_cache(self) -> WanInferencePipelineCache: """Initialize the autoregressive cache.""" @@ -151,13 +149,9 @@ def run(self) -> None: return # Write the video. - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "t c h w -> t h w c") - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=config.fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -166,8 +160,9 @@ def run(self) -> None: # Write the perf stats. if stats_history: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + config.output_dir, config.runner_name, stats_history + ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) diff --git a/integrations/flashvsr/flashvsr/runner.py b/integrations/flashvsr/flashvsr/runner.py index fcd6be577..5894cd73f 100644 --- a/integrations/flashvsr/flashvsr/runner.py +++ b/integrations/flashvsr/flashvsr/runner.py @@ -17,23 +17,28 @@ from __future__ import annotations -import json import os from dataclasses import dataclass, field from pathlib import Path from typing import Literal -import mediapy as media -import numpy as np import torch -from einops import rearrange from loguru import logger from flashdreams.core.distributed import init as init_distributed -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.config import derive_config from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig, _is_torchrun_env +from flashdreams.infra.runner_io import ( + ensure_output_dir, + read_video_fps, + read_video_rgb, + resolve_input_path, + rgb_video_to_normalized_tensor, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashvsr.encoder import FlashVSREncoder from flashvsr.pipeline import ( FlashVSRPipeline, @@ -59,21 +64,6 @@ """User-writable cache for on-the-fly demo-video downloads.""" -def _resolve_input_path(input_path: str | Path) -> Path: - """Return a local ``Path`` for ``input_path``, downloading URLs on the fly. - - ``http(s)://`` strings are atomically fetched into - :data:`INPUT_CACHE_DIR` and validated as decodable videos before being - published; local paths pass through unchanged. - """ - if isinstance(input_path, Path): - return input_path - if not input_path.startswith(("http://", "https://")): - return Path(input_path) - - return download_to_cache(input_path, cache_dir=INPUT_CACHE_DIR) - - ## Chunk planning helpers @@ -184,10 +174,7 @@ def _probe_input_fps(path: Path) -> float: Probed frame rate, or ``30.0`` if metadata probing fails. """ try: - # ``mediapy`` ships without type stubs so ty cannot see the - # ``VideoMetadata.from_path`` classmethod (added in mediapy 1.1). - meta = media.VideoMetadata.from_path(str(path)) # ty: ignore[unresolved-attribute] - return float(meta.fps) + return read_video_fps(path) except Exception: logger.warning(f"Could not probe fps for {path}; defaulting to 30.0.") return 30.0 @@ -327,7 +314,7 @@ def _load_input_video(self) -> tuple[torch.Tensor, int, int, float]: # Resolve once: local paths pass through, ``http(s)://`` URLs are # downloaded into :data:`INPUT_CACHE_DIR` and validated as # decodable videos before being published. - path = _resolve_input_path(config.input_path) + path = resolve_input_path(config.input_path, cache_dir=INPUT_CACHE_DIR) assert path.is_file(), ( f"--input-path must resolve to an existing video file; got: " f"{config.input_path!r} -> {path!r}" @@ -335,7 +322,7 @@ def _load_input_video(self) -> tuple[torch.Tensor, int, int, float]: if self.is_rank_zero: logger.info(f"Reading {path} ...") - video_np = media.read_video(str(path))[..., :3] # uint8 [T, H, W, C] + video_np = read_video_rgb(path) T, H, W, _ = video_np.shape if config.crop_region != "none": @@ -354,8 +341,12 @@ def _load_input_video(self) -> tuple[torch.Tensor, int, int, float]: # bf16/cuda placement is deferred to :meth:`run` once the # pipeline (and therefore the target device + dtype) exists. - video_t = torch.from_numpy(video_np.astype(np.float32)) / 127.5 - 1.0 - video_t = rearrange(video_t, "T H W C -> 1 C T H W") + video_t = rgb_video_to_normalized_tensor( + video_np, + device=torch.device("cpu"), + dtype=torch.float32, + ) + video_t = video_t.permute(1, 0, 2, 3).unsqueeze(0) return video_t, H, W, fps def _initialize_cache(self) -> FlashVSRPipelineCache: @@ -470,12 +461,9 @@ def run(self) -> None: if generated is None: return - # [1, 3, T_out, target_H, target_W] in [-1, 1] -> [T_out, H, W, 3]. - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "1 c t h w -> t h w c") - arr = ((canvas.float().numpy() + 1.0) / 2.0 * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=fps, layout="bcthw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -483,8 +471,9 @@ def run(self) -> None: ) if stats_history: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + config.output_dir, config.runner_name, stats_history + ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> " f"{stats_path.resolve()}" diff --git a/integrations/hy_worldplay/hy_worldplay/runner.py b/integrations/hy_worldplay/hy_worldplay/runner.py index 84865e562..7999d4371 100644 --- a/integrations/hy_worldplay/hy_worldplay/runner.py +++ b/integrations/hy_worldplay/hy_worldplay/runner.py @@ -18,7 +18,6 @@ from __future__ import annotations import contextlib -import json import time from dataclasses import dataclass, field from pathlib import Path @@ -31,6 +30,12 @@ from flashdreams.core.io.download import download_to_cache from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, +) from flashdreams.recipes.wan.pipeline import WanInferencePipeline __all__ = [ @@ -121,16 +126,6 @@ def _pil_to_numpy(img: object) -> object: return np.asarray(img) -def _resolve_prompt(value: str | Path) -> str: - """Read an inline prompt or the first non-empty line of a prompt file.""" - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value - - def _write_mp4(video: Tensor, out_path: Path, *, fps: int) -> None: """Persist a decoded video tensor as mp4. @@ -351,7 +346,7 @@ def run(self) -> None: image = preprocess_first_frame( cfg.image_path, cfg.pixel_height, cfg.pixel_width ).to(device=device, dtype=first_param.dtype) - prompt = _resolve_prompt(cfg.prompt) + prompt = resolve_prompt_value(cfg.prompt) cache = self.pipeline.initialize_cache( text=[prompt], @@ -406,8 +401,8 @@ def run(self) -> None: if video is None: return - cfg.output_dir.mkdir(parents=True, exist_ok=True) - out_path = cfg.output_dir / f"{cfg.runner_name}.mp4" + ensure_output_dir(cfg.output_dir) + out_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") _write_mp4(video, out_path, fps=cfg.fps) logger.info( f"[{cfg.runner_name}] wrote video " @@ -415,8 +410,9 @@ def run(self) -> None: ) if stats_history: - stats_path = cfg.output_dir / f"stats_{cfg.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + cfg.output_dir, cfg.runner_name, stats_history + ) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) diff --git a/integrations/lingbot/lingbot/runner.py b/integrations/lingbot/lingbot/runner.py index cf585806e..14efa50fe 100644 --- a/integrations/lingbot/lingbot/runner.py +++ b/integrations/lingbot/lingbot/runner.py @@ -17,19 +17,24 @@ from __future__ import annotations -import json from dataclasses import dataclass, field from pathlib import Path import numpy as np import torch -from einops import rearrange from loguru import logger from flashdreams.core.io.disk import default_flashdreams_cache_dir from flashdreams.core.io.download import download_to_cache from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from lingbot.encoder.camctrl import CamCtrlInput from lingbot.encoder.utils import ( get_Ks_transformed, @@ -240,11 +245,14 @@ def run(self) -> None: # shipped configs pin ``batch_shape=()`` so a single-rollout layout # is just ``[T, C, H, W]`` (image) / ``[T, 4, 4]`` (poses) / # ``[T, 4]`` (intrinsics). - first_frames_t = _load_first_frame( + first_frames_t = load_first_frame_tensor( cfg.image_path, pixel_height=cfg.pixel_height, pixel_width=cfg.pixel_width, device=device, + dtype=torch.bfloat16, + interpolation="cubic", + install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", ) Ks = np.load(cfg.intrinsic_path) @@ -318,62 +326,24 @@ def run(self) -> None: if video is None: return - cfg.output_dir.mkdir(parents=True, exist_ok=True) - canvas = rearrange(video, "t c h w -> t h w c") - video_path = cfg.output_dir / f"{cfg.runner_name}.mp4" - _write_video(canvas, video_path, fps=cfg.fps) + ensure_output_dir(cfg.output_dir) + video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") + write_video_tensor( + video, + video_path, + fps=cfg.fps, + layout="tchw", + install_hint="Install the lingbot plugin: pip install flashdreams-lingbot.", + ) logger.info( f"[{cfg.runner_name}] wrote video {tuple(video.shape)} " f"-> {video_path.resolve()}" ) if stats_history: - stats_path = cfg.output_dir / f"stats_{cfg.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + cfg.output_dir, cfg.runner_name, stats_history + ) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) - - -## I/O helpers (``cv2`` / ``mediapy`` are listed under ``flashdreams-lingbot``'s -## runtime dependencies, so the import-time guards mostly catch the bare -## ``pip install flashdreams`` case where the plugin extras were skipped). - - -def _load_first_frame( - path: Path, *, pixel_height: int, pixel_width: int, device: torch.device -) -> torch.Tensor: - """Load + resize a first-frame image into ``[1, C, H, W]`` ``[-1, 1]``.""" - try: - import cv2 # noqa: PLC0415 - import mediapy as media # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Loading the first-frame image needs mediapy + opencv. " - "Install the lingbot plugin: pip install flashdreams-lingbot." - ) from exc - - arr = media.read_image(str(path))[..., :3] - # Bicubic to match the upstream Lingbot World demo / generate_fast.py - # (which uses ``F.interpolate(mode='bicubic')`` over the ``[-1, 1]`` - # tensor); bilinear here would give a different first-frame VAE latent. - arr = cv2.resize(arr, (pixel_width, pixel_height), interpolation=cv2.INTER_CUBIC) - tensor = ( - torch.from_numpy(arr).to(device=device, dtype=torch.bfloat16) / 127.5 - 1.0 - ) # [H, W, 3] - return rearrange(tensor, "h w c -> 1 c h w") # [T=1, C, H, W] - - -def _write_video(canvas: torch.Tensor, path: Path, *, fps: int) -> None: - """Save a ``[T, H, W, C]`` ``[-1, 1]`` tensor as an MP4.""" - try: - import mediapy as media # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Writing the output video needs mediapy. " - "Install the lingbot plugin: pip install flashdreams-lingbot." - ) from exc - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(path), arr, fps=fps) diff --git a/integrations/omnidreams/omnidreams/runner.py b/integrations/omnidreams/omnidreams/runner.py index b2cd76148..b0322e11b 100644 --- a/integrations/omnidreams/omnidreams/runner.py +++ b/integrations/omnidreams/omnidreams/runner.py @@ -27,11 +27,9 @@ from __future__ import annotations -import json from dataclasses import dataclass, field from pathlib import Path -import numpy as np import torch from einops import rearrange from loguru import logger @@ -43,6 +41,15 @@ from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + DEFAULT_RUNNER_INSTALL_HINT, + ensure_output_dir, + load_first_frame_tensor, + load_video_tensor, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) DEFAULT_VIDEO_HEIGHT = 704 """Pixel-space rollout height (matches the trained 720p chassis).""" @@ -50,8 +57,6 @@ DEFAULT_VIDEO_WIDTH = 1280 """Pixel-space rollout width (matches the trained 720p chassis).""" -IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".webp"} - _REPO_ROOT = Path(__file__).resolve().parents[4] EXAMPLE_DATA_HF_REPO = "nvidia/omni-dreams-samples" @@ -311,7 +316,7 @@ def _run_save_embeddings(self, output_path: Path) -> None: text_emb = embeddings["text_embeddings"] image_emb = embeddings["image_embeddings"] assert text_emb is not None and image_emb is not None - output_path.parent.mkdir(parents=True, exist_ok=True) + ensure_output_dir(output_path.parent) torch.save(embeddings, output_path) logger.info( f"[{cfg.runner_name}] saved precomputed embeddings " @@ -432,9 +437,15 @@ def _rollout_and_save( video=video, ) - cfg.output_dir.mkdir(parents=True, exist_ok=True) - video_path = cfg.output_dir / f"{cfg.runner_name}.mp4" - _write_video(canvas, video_path, fps=cfg.output_fps) + ensure_output_dir(cfg.output_dir) + video_path = runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4") + write_video_tensor( + canvas, + video_path, + fps=cfg.output_fps, + layout="thwc", + install_hint=DEFAULT_RUNNER_INSTALL_HINT, + ) logger.info( f"[{cfg.runner_name}] wrote {output_description} {tuple(canvas.shape)} " f"from generated={tuple(video.shape)}; wrote video {tuple(video.shape)} " @@ -442,8 +453,9 @@ def _rollout_and_save( ) if stats_history: - stats_path = cfg.output_dir / f"stats_{cfg.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + cfg.output_dir, cfg.runner_name, stats_history + ) logger.info( f"[{cfg.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) @@ -490,12 +502,14 @@ def _load_first_frames( """Load the per-camera first-frame seeds as ``[B=1, V, 1, C, H, W]``.""" cfg = self.config first_frames = [ - _load_first_frame( + load_first_frame_tensor( p, pixel_height=cfg.pixel_height, pixel_width=cfg.pixel_width, device=device, dtype=dtype, + allow_video=True, + install_hint=DEFAULT_RUNNER_INSTALL_HINT, ) for p in first_frame_paths ] @@ -549,51 +563,6 @@ def _resolve_paths( ] -## I/O helpers (``cv2`` / ``mediapy`` lazy-imported; live under the ``runners`` extras). - - -def _read_first_frame_np(path: Path) -> np.ndarray: - """Read a first-frame image (or frame 0 of a video) as ``[H, W, 3]``.""" - try: - import mediapy as media # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Loading the first-frame asset needs mediapy. " - "Install the runner extras: pip install 'flashdreams[runners]'." - ) from exc - - if path.suffix.lower() in IMAGE_SUFFIXES: - return media.read_image(str(path))[..., :3] - video = media.read_video(str(path)) - assert video.shape[0] > 0, f"video has no frames: {path}" - return video[0, ..., :3] - - -def _load_first_frame( - path: Path, - *, - pixel_height: int, - pixel_width: int, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - """Resize a first-frame asset and return ``[1, C, H, W]`` in ``[-1, 1]``.""" - try: - import cv2 # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Resizing the first-frame asset needs opencv. " - "Install the runner extras: pip install 'flashdreams[runners]'." - ) from exc - - arr = _read_first_frame_np(path) - arr = cv2.resize(arr, (pixel_width, pixel_height)) - tensor = ( - torch.from_numpy(arr).to(dtype=dtype, device=device) / 127.5 - 1.0 - ) # [H, W, C] - return rearrange(tensor, "h w c -> 1 c h w") - - def _load_video( path: Path, *, @@ -603,36 +572,11 @@ def _load_video( dtype: torch.dtype, ) -> torch.Tensor: """Load + resize an HDMap video to ``[T, C, H, W]`` in ``[-1, 1]``.""" - try: - import cv2 # noqa: PLC0415 - import mediapy as media # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Loading HDMap videos needs mediapy + opencv. " - "Install the runner extras: pip install 'flashdreams[runners]'." - ) from exc - - video_np = media.read_video(str(path))[..., :3] - if video_np.shape[1:3] != (pixel_height, pixel_width): - video_np = np.stack( - [cv2.resize(f, (pixel_width, pixel_height)) for f in video_np], axis=0 - ) - tensor = ( - torch.from_numpy(video_np).to(dtype=dtype, device=device) / 127.5 - 1.0 - ) # [T, H, W, C] - return rearrange(tensor, "t h w c -> t c h w") - - -def _write_video(canvas: torch.Tensor, path: Path, *, fps: int) -> None: - """Save a ``[T, H, W, C]`` ``[-1, 1]`` tensor as an MP4.""" - try: - import mediapy as media # noqa: PLC0415 - except ImportError as exc: # pragma: no cover - import-time gate - raise ImportError( - "Writing the output video needs mediapy. Install the runner " - "extras: pip install 'flashdreams[runners]'." - ) from exc - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(path), arr, fps=fps) + return load_video_tensor( + path, + pixel_height=pixel_height, + pixel_width=pixel_width, + device=device, + dtype=dtype, + install_hint=DEFAULT_RUNNER_INSTALL_HINT, + ) diff --git a/integrations/self_forcing/self_forcing/runner.py b/integrations/self_forcing/self_forcing/runner.py index 754f742f0..6a04c3b0f 100644 --- a/integrations/self_forcing/self_forcing/runner.py +++ b/integrations/self_forcing/self_forcing/runner.py @@ -17,17 +17,21 @@ from __future__ import annotations -import json from dataclasses import dataclass, field from pathlib import Path -import mediapy as media -from einops import rearrange from loguru import logger from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashdreams.recipes.wan import ( WanInferencePipeline, WanInferencePipelineCache, @@ -91,13 +95,7 @@ def _resolve_prompt(self) -> str: A Path reads its first non-empty line, a str is used as-is. """ - value = self.config.prompt - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value + return resolve_prompt_value(self.config.prompt) def _initialize_cache(self) -> WanInferencePipelineCache: """Initialize the autoregressive cache.""" @@ -149,13 +147,9 @@ def run(self) -> None: return # Write the video. - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "t c h w -> t h w c") - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=config.fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -164,8 +158,9 @@ def run(self) -> None: # Write the perf stats. if stats_history: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text(json.dumps(stats_history, indent=2)) + stats_path = write_runner_stats( + config.output_dir, config.runner_name, stats_history + ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" ) diff --git a/integrations/wan21/wan21/runner.py b/integrations/wan21/wan21/runner.py index 23a44dbe9..9bfb04607 100644 --- a/integrations/wan21/wan21/runner.py +++ b/integrations/wan21/wan21/runner.py @@ -17,21 +17,26 @@ from __future__ import annotations -import json import os from dataclasses import dataclass, field from pathlib import Path -import cv2 -import mediapy as media import torch -from einops import rearrange from loguru import logger -from flashdreams.core.io.download import download_to_cache from flashdreams.infra.decoder import StreamingVideoDecoder from flashdreams.infra.postprocess import VideoTensorLayout from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ( + ensure_output_dir, + load_first_frame_tensor, + read_image_rgb, + resolve_input_path, + resolve_prompt_value, + runner_artifact_path, + write_runner_stats, + write_video_tensor, +) from flashdreams.recipes.wan import ( WanInferencePipeline, WanInferencePipelineCache, @@ -67,25 +72,6 @@ """User-writable cache for on-the-fly I2V first-frame downloads.""" -def _resolve_image_path(image_path: str | Path) -> Path: - """Return a local ``Path`` for ``image_path``, downloading URLs on the fly. - - ``http(s)://`` strings are atomically fetched into - :data:`IMAGE_CACHE_DIR` and validated as decodable images before - being published; local paths pass through unchanged. - """ - if isinstance(image_path, Path): - return image_path - if not image_path.startswith(("http://", "https://")): - return Path(image_path) - - return download_to_cache( - image_path, - cache_dir=IMAGE_CACHE_DIR, - validator=lambda p: media.read_image(str(p)), - ) - - @dataclass(kw_only=True) class Wan21T2VRunnerConfig(RunnerConfig): """Runner config for the Wan 2.1 T2V variant. @@ -156,13 +142,7 @@ def _resolve_prompt(self) -> str: A Path reads its first non-empty line, a str is used as-is. """ - value = self.config.prompt - if isinstance(value, Path): - lines = [ln.strip() for ln in value.read_text().splitlines() if ln.strip()] - assert lines, f"prompt file {value} has no non-empty lines" - return lines[0] - assert value, "--prompt must be a non-empty string or a path to a .txt file" - return value + return resolve_prompt_value(self.config.prompt) def _initialize_cache(self) -> WanInferencePipelineCache: """Initialize the autoregressive cache for T2V.""" @@ -201,13 +181,9 @@ def run(self) -> None: return # Write the video. - config.output_dir.mkdir(parents=True, exist_ok=True) - video_path = config.output_dir / f"{config.runner_name}.mp4" - canvas = rearrange(generated, "t c h w -> t h w c") - - arr = (canvas.float().numpy() + 1.0) / 2.0 - arr = (arr * 255).clip(0, 255).astype("uint8") - media.write_video(str(video_path), arr, fps=config.fps) + ensure_output_dir(config.output_dir) + video_path = runner_artifact_path(config.output_dir, config.runner_name, "mp4") + write_video_tensor(generated, video_path, fps=config.fps, layout="tchw") logger.info( f"[{config.runner_name}] wrote video {tuple(generated.shape)} " @@ -216,9 +192,10 @@ def run(self) -> None: # Write the perf stats. if stats is not None: - stats_path = config.output_dir / f"stats_{config.runner_name}.json" - stats_path.write_text( - json.dumps([{"autoregressive_index": 0, **stats}], indent=2) + stats_path = write_runner_stats( + config.output_dir, + config.runner_name, + [{"autoregressive_index": 0, **stats}], ) logger.info( f"[{config.runner_name}] wrote per-AR-step stats -> {stats_path.resolve()}" @@ -248,14 +225,16 @@ def _initialize_cache(self) -> WanInferencePipelineCache: # in shape [T=1, C, H, W] (matches batch_shape=()). Pin to the # pipeline's actual device so non-default ``--device`` selections # (and the auto cuda:LOCAL_RANK override under torchrun) both work. - image_path = _resolve_image_path(config.image_path) - arr = media.read_image(str(image_path))[..., :3] - arr = cv2.resize(arr, (config.pixel_width, config.pixel_height)) - tensor = ( - torch.from_numpy(arr).to(device=self.pipeline.device, dtype=torch.bfloat16) - / 127.5 - - 1.0 + image = load_first_frame_tensor( + resolve_input_path( + config.image_path, + cache_dir=IMAGE_CACHE_DIR, + validator=read_image_rgb, + ), + pixel_height=config.pixel_height, + pixel_width=config.pixel_width, + device=self.pipeline.device, + dtype=torch.bfloat16, ) - image = rearrange(tensor, "h w c -> 1 c h w") return self.pipeline.initialize_cache(text=[prompt], image=image)