From 952ba6155464c6f90ced515594811ff08e051ed6 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Tue, 7 Jul 2026 03:54:22 +0000 Subject: [PATCH 1/5] feat: Cosmos3 videogen worker on public cosmos-framework (overlay) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Serving overlay for nvidia/Cosmos3-Nano and Cosmos3-Super on NVIDIA's public cosmos-framework, replacing the retired private EA fork (deepinfra/cosmos3): - worker.py: ported from the EA fork (import renames only); dynamo registration, schemas, gloo control-plane fix, torchrun relaunch and standalone mode unchanged. - Dockerfile: overlay build — clones cosmos-framework at a pinned ref, applies patches/, mirrors upstream's build, bakes worker + HF prefetch (incl. gated Cosmos-Guardrail1 for default-on guardrails). - patches/0001: encoder -threads cap (EA PR #19 re-applied). The 16 EA-era weight/code mismatch shims are obsolete upstream (native diffusers loading, AVAE released, parallelism-config bug fixed, guardrails default on). See examples/cosmos3/README.md for provenance. Co-Authored-By: Claude Fable 5 --- examples/cosmos3/Dockerfile | 133 ++++ examples/cosmos3/README.md | 39 ++ .../patches/0001-cap-encoder-threads.patch | 33 + examples/cosmos3/worker.py | 645 ++++++++++++++++++ 4 files changed, 850 insertions(+) create mode 100644 examples/cosmos3/Dockerfile create mode 100644 examples/cosmos3/README.md create mode 100644 examples/cosmos3/patches/0001-cap-encoder-threads.patch create mode 100644 examples/cosmos3/worker.py diff --git a/examples/cosmos3/Dockerfile b/examples/cosmos3/Dockerfile new file mode 100644 index 000000000000..706778f6bb8a --- /dev/null +++ b/examples/cosmos3/Dockerfile @@ -0,0 +1,133 @@ +# DeepInfra Cosmos3 Dynamo videogen worker image. +# +# Overlay build: clones NVIDIA/cosmos-framework at a pinned ref, applies the +# reviewed patches from ./patches, then follows upstream's own Dockerfile steps +# (same base image, uv-managed env) and bakes in the DeepInfra worker. +# There is no DeepInfra copy of the framework repo — this directory plus the +# pin below is the complete, reviewable definition of the image. +# +# Build (from this directory): +# docker build --build-arg HF_TOKEN= \ +# -t localhost:30500/cosmos3-worker:v1 . + +ARG CUDA_VERSION=13.0.2 +ARG BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04 +FROM ${BASE_IMAGE} + +ENV DEBIAN_FRONTEND=noninteractive + +RUN --mount=type=cache,target=/var/cache/apt \ + --mount=type=cache,target=/var/lib/apt \ + apt-get update && \ + apt-get install -y --no-install-recommends \ + curl \ + ffmpeg \ + git \ + git-lfs \ + tree \ + wget + +COPY --from=ghcr.io/astral-sh/uv:0.10.8 /uv /uvx /usr/local/bin/ +ENV UV_LINK_MODE=copy +ENV UV_PYTHON_CACHE_DIR=/root/.cache/uv/python +ENV PATH="/root/.local/bin:$PATH" + +# Upstream pin. Bumping this ref is THE upstream-update mechanism; if a patch +# in ./patches stops applying, the build fails loudly and the patch must be +# rebased (or dropped because upstream absorbed the fix). +ARG COSMOS_FRAMEWORK_REPO=https://github.com/NVIDIA/cosmos-framework.git +ARG COSMOS_FRAMEWORK_REF=463ac142b47c637098a7e86d0631252f82b65418 + +WORKDIR /workspace +RUN git clone --filter=blob:none ${COSMOS_FRAMEWORK_REPO} /workspace && \ + git -C /workspace checkout ${COSMOS_FRAMEWORK_REF} + +# DeepInfra patches to upstream code (reviewed in deepinfra/dynamo). +COPY patches /workspace/deepinfra-patches +RUN git -C /workspace apply --verbose deepinfra-patches/*.patch + +# Upstream build steps (mirrors NVIDIA/cosmos-framework Dockerfile @ the pin, +# minus bind-mounts: the sources are already in the layer from the clone). +RUN --mount=type=cache,target=/root/.cache/uv uv python install +RUN echo "$CUDA_VERSION" | sed -E 's/^([0-9]+)\.([0-9]+).*/cu\1\2/' > /root/.cuda-name +RUN --mount=type=cache,target=/root/.cache/uv \ + uv sync --locked --no-install-project --no-editable --all-extras --group=$(cat /root/.cuda-name) +ENV PATH="/workspace/.venv/bin:$PATH" + +# Triton bundled ptxas doesn't support latest GPU architectures +ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas" + +# Serving deps. fastapi/uvicorn for standalone mode; ai-dynamo + uvloop for +# production Dynamo mode. ai-dynamo provides the DistributedRuntime, frontend +# binary, and endpoint registration used by the videogen worker pattern. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install fastapi uvicorn ai-dynamo==1.1.1 uvloop + +# Pre-install the HF CLI tool so that `uvx hf@1.13.0` reuses this +# installation instead of creating an ephemeral env missing `click`. +RUN --mount=type=cache,target=/root/.cache/uv \ + uv tool install hf@1.13.0 --with click + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --no-deps -e /workspace + +# Pre-fetch HuggingFace dependencies needed by cosmos_framework at model init +# time: Qwen tokenizer configs (*.json, *.txt — few MB), Wan VAE checkpoints, +# and the full Cosmos-Guardrail1 snapshot (gated repo: the build HF_TOKEN's +# account must have accepted access; gating is auto-approve). +# This eliminates any runtime dependency on HuggingFace availability. +# Revisions MUST match the pins in +# cosmos_framework/inference/common/checkpoints.py and +# cosmos_framework/auxiliary/guardrail/common/core.py at the ref above — the +# runtime looks up these exact revisions; a mismatch silently falls back to a +# live (unauthenticated, rate-limited) HF fetch at every pod boot. +ARG HF_TOKEN="" +RUN HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Qwen/Qwen3-0.6B \ + --repo-type model \ + --revision c1899de289a04d12100db370d81485cdf75e47ca \ + --include '*.json' --include '*.txt' && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Qwen/Qwen3-VL-2B-Instruct \ + --repo-type model \ + --revision 89644892e4d85e24eaac8bacfd4f463576704203 \ + --include '*.json' --include '*.txt' && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Qwen/Qwen3-VL-8B-Instruct \ + --repo-type model \ + --revision 0c351dd01ed87e9c1b53cbc748cba10e6187ff3b \ + --include '*.json' --include '*.txt' && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Qwen/Qwen3-VL-32B-Instruct \ + --repo-type model \ + --revision 0cfaf48183f594c314753d30a4c4974bc75f3ccb \ + --include '*.json' --include '*.txt' && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Wan-AI/Wan2.1-T2V-14B \ + --repo-type model \ + --revision a064a6c71f5be440641209c07bf2a5ce7a2ff5e4 \ + Wan2.1_VAE.pth && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + Wan-AI/Wan2.2-TI2V-5B \ + --repo-type model \ + --revision 921dbaf3f1674a56f47e83fb80a34bac8a8f203e \ + Wan2.2_VAE.pth && \ + HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + nvidia/Cosmos-Guardrail1 \ + --repo-type model \ + --revision d6d4bfa899a71454a700907664f3e88f503950cf + +# DeepInfra runtime settings. +# - LD_LIBRARY_PATH is cleared so the venv/CUDA libs resolve without host leakage. +# - torch.compile (inductor) and triton write kernel caches to a persistent +# mountable location instead of an ephemeral home dir. +ENV LD_LIBRARY_PATH="" +ENV TORCHINDUCTOR_CACHE_DIR="/cache/torchinductor" +ENV TRITON_CACHE_DIR="/cache/triton" +RUN mkdir -p /cache/torchinductor /cache/triton + +# DeepInfra Dynamo videogen worker. Runtime args +# (--served-model-name, --model, --num-gpus, ...) are appended to this command. +# worker.py self-relaunches under torchrun when --num-gpus > 1. +COPY worker.py /workspace/worker.py +ENTRYPOINT ["python3", "worker.py"] diff --git a/examples/cosmos3/README.md b/examples/cosmos3/README.md new file mode 100644 index 000000000000..e07eb6ffaa00 --- /dev/null +++ b/examples/cosmos3/README.md @@ -0,0 +1,39 @@ +# Cosmos3 Dynamo videogen worker (public cosmos-framework) + +DeepInfra serving overlay for `nvidia/Cosmos3-Nano` and `nvidia/Cosmos3-Super` +on NVIDIA's **public** [cosmos-framework](https://github.com/NVIDIA/cosmos-framework), +replacing the retired private EA fork (`deepinfra/cosmos3`, archived — see +provenance below). + +## What lives here + +- `worker.py` — the Dynamo videogen worker: registers with the Dynamo frontend + (`ModelType.Videos`), translates deepapi requests to + `cosmos_framework.inference` calls, multi-GPU via torchrun with a gloo + control-plane group (idle NCCL-watchdog fix). `--standalone` runs a plain + FastAPI server for testing without any Dynamo infra. +- `Dockerfile` — overlay build: clones cosmos-framework at the pinned + `COSMOS_FRAMEWORK_REF`, applies `patches/`, runs upstream's own build steps, + bakes worker + HF prefetch (incl. gated `nvidia/Cosmos-Guardrail1`; the build + `HF_TOKEN` account must have accepted that repo's auto-approve gate). +- `patches/` — the complete set of in-tree changes to upstream: + - `0001-cap-encoder-threads.patch` — cap libx264 `-threads` (default 32, + `DI_VIDEO_ENCODE_THREADS` to override) so mp4 encode doesn't grab every + core on shared multi-GPU nodes. + +## Updating upstream + +Bump `COSMOS_FRAMEWORK_REF` in the Dockerfile and rebuild. If a patch fails to +apply, the build fails: rebase the patch, or delete it if upstream absorbed +the fix. + +## Provenance + +The EA-era history (19 PRs) lives in the archived private repo +`deepinfra/cosmos3`. Sixteen of those PRs were shims to load the public +diffusers weights on the private EA code — all obsolete: the public framework +loads the public checkpoints natively, ships the audio (AVAE) tokenizer, +fixes the discarded-parallelism-config bug (`inference/model.py`, fixed +upstream), and enables guardrails by default. What survives is `worker.py` +(EA PRs #1/#18), the encoder-thread cap (EA PR #19, now `patches/0001`), and +the Dockerfile's prefetch/bake strategy. diff --git a/examples/cosmos3/patches/0001-cap-encoder-threads.patch b/examples/cosmos3/patches/0001-cap-encoder-threads.patch new file mode 100644 index 000000000000..3cd17fcc7b58 --- /dev/null +++ b/examples/cosmos3/patches/0001-cap-encoder-threads.patch @@ -0,0 +1,33 @@ +diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +index fdc944f..8393fa5 100644 +--- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py ++++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +@@ -1,6 +1,7 @@ + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + # SPDX-License-Identifier: OpenMDW-1.1 + ++import os + from typing import IO, Any, Dict, Tuple + + import imageio +@@ -171,6 +172,20 @@ class ImageioVideoHandler(BaseFileHandler): + } + # Update with any other kwargs + mimsave_kwargs.update(kwargs) ++ ++ # deepinfra: cap CPU encoder threads. imageio encodes mp4 with CPU ++ # libx264; with no -threads, x264 auto-detects every core and ++ # oversubscribes on shared multi-GPU nodes (e.g. 8 GPUs / 224 cores), ++ # contending with neighbouring GPU pods so encode latency drifts with ++ # node load. Append a per-GPU-sized cap unless a caller already ++ # specified -threads. Override via DI_VIDEO_ENCODE_THREADS. ++ if "-threads" not in mimsave_kwargs.get("ffmpeg_params", []): ++ mimsave_kwargs["ffmpeg_params"] = [ ++ *mimsave_kwargs.get("ffmpeg_params", []), ++ "-threads", ++ os.getenv("DI_VIDEO_ENCODE_THREADS", "32"), ++ ] ++ + log.debug(f"mimsave_kwargs: {mimsave_kwargs}") + + imageio.mimsave(file, obj, format, **mimsave_kwargs) diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py new file mode 100644 index 000000000000..daa76fcf0d6f --- /dev/null +++ b/examples/cosmos3/worker.py @@ -0,0 +1,645 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +"""Cosmos 3 Dynamo videogen worker. + +Loads a Cosmos 3 checkpoint once at startup and serves text/image-to-video +generation. Two operating modes: + +**Dynamo mode** (production, default): + Registers as a Dynamo backend endpoint. A separate Dynamo frontend + (``python -m dynamo.frontend``) handles HTTP and routes requests to + this worker via Dynamo RPC. The frontend exposes ``/v1/videos`` and + this worker registers with ``ModelType.Videos``. + +**Standalone mode** (``--standalone``, for local testing): + Runs a plain FastAPI HTTP server directly on the worker, no frontend + needed. Useful for ``docker run`` one-shot testing. + +In both modes the underlying inference is identical: + + setup_args = OmniSetupOverrides(...).build_setup(world_size=WORLD_SIZE) + pipe = OmniInference.create(setup_args) + sample_args = OmniSampleOverrides(...).build_sample(model_config=pipe.model_config) + pipe.generate([sample_args]) + +Multi-GPU (CFG- and context-parallel) inference: + +``pipe.generate`` is a *collective* — every rank must call it with the same +sample args. For ``--num-gpus > 1`` the worker (re-)launches itself under +``torchrun --nproc-per-node=N``; ``init_script()`` then initializes the process +group. Rank 0 runs the endpoint and, for each request, broadcasts the request +to all ranks; every rank runs ``generate`` together, and only rank 0 reads back +the produced media file. + +Usage: + # Dynamo mode (requires a frontend running alongside): + python3 worker.py --served-model-name nvidia/Cosmos3-Nano --model Cosmos3-Nano + + # Standalone HTTP mode (no frontend needed): + python3 worker.py --standalone --served-model-name nvidia/Cosmos3-Nano --model Cosmos3-Nano +""" + +# init_script() configures logging/CUDA and, when WORLD_SIZE > 1 (i.e. under +# torchrun), initializes torch.distributed. It must run before any heavy +# cosmos_framework/torch imports, exactly as in cosmos_framework/scripts/inference.py. +from cosmos_framework.inference.common.init import init_script + +init_script() + +import argparse +import asyncio +import base64 +import mimetypes +import os +import random +import re +import shutil +import sys +import threading +import time +import uuid +from pathlib import Path +from typing import Any + +import torch +from pydantic import BaseModel, Field + +from cosmos_framework.inference.args import OmniSampleOverrides, OmniSetupOverrides +from cosmos_framework.inference.common.init import get_local_rank, get_rank, init_output_dir, is_rank0 +from cosmos_framework.inference.inference import OmniInference +from cosmos_framework.utils import distributed, log + +# --------------------------------------------------------------------------- +# Control-plane process group (multi-GPU request dispatch) +# --------------------------------------------------------------------------- +# The default process group is NCCL, whose watchdog aborts any collective +# pending longer than the init timeout (default 30 min). The idle +# _worker_rank_loop blocks on broadcast_object() waiting for the *next* +# request, so a 2-GPU worker that receives no traffic for 30 minutes is +# killed by its own watchdog (SeqNum stuck at the first idle wait; pager +# incident 2026-06-07). Request dispatch therefore goes over a gloo (CPU) +# group with a generous timeout; NCCL remains in charge of the actual +# generation collectives inside _run_generate(). +# +# new_group() is itself collective: every rank executes this at import time +# (module import is symmetric under torchrun), keeping creation in lockstep. +_CONTROL_PG = None +if torch.distributed.is_available() and torch.distributed.is_initialized(): + import datetime as _dt + + _CONTROL_PG = torch.distributed.new_group( + backend="gloo", timeout=_dt.timedelta(days=7)) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_RESOLUTION_ALIASES = {"256p": "256", "480p": "480", "720p": "720", "1080p": "1080"} +_DATA_URI_RE = re.compile(r"^data:([^;,]+)?(?:;base64)?,(.*)$", re.DOTALL) + +# Reverse lookup: (width, height) -> (resolution, aspect_ratio). +# Built from cosmos_framework's VIDEO_RES_SIZE_INFO (same as IMAGE_RES_SIZE_INFO +# for resolutions we support). Used to decode the Dynamo frontend's "size" field +# ("WxH") back to the resolution + aspect_ratio the inference engine expects. +_SIZE_TO_RES_AR: dict[tuple[int, int], tuple[str, str]] = {} +def _init_size_lookup() -> None: + from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO + for res, ar_map in VIDEO_RES_SIZE_INFO.items(): + for ar, (w, h) in ar_map.items(): + _SIZE_TO_RES_AR[(w, h)] = (res, ar) +_init_size_lookup() + + +def _normalize_resolution(resolution: str) -> str: + return _RESOLUTION_ALIASES.get(resolution.strip().lower(), resolution.strip()) + + +def _decode_data_uri(uri: str, output_dir: Path) -> str: + """Decode a ``data:`` URI to a local file and return the path. + + Cosmos 3's ``vision_path`` only accepts ``http(s)://``, ``s3://``, or + local file paths. This helper bridges the gap for base64-encoded payloads + sent by the API gateway. + """ + m = _DATA_URI_RE.match(uri) + if m is None: + raise ValueError(f"Invalid data URI: {uri[:80]}...") + mime_type = m.group(1) or "application/octet-stream" + ext = mimetypes.guess_extension(mime_type) or ".bin" + data = base64.b64decode(m.group(2)) + path = output_dir / f"input_vision{ext}" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return str(path) + + +# --------------------------------------------------------------------------- +# Request / Response models +# --------------------------------------------------------------------------- +# These match the JSON format that deepinfra/deepapi/handler.py sends to +# the /v1/videos endpoint and the response it expects back. + + +# Valid model_mode values for action modalities. +_ACTION_MODES = frozenset({"forward_dynamics", "inverse_dynamics", "policy"}) + + +class NvExtFields(BaseModel): + """Subset of Dynamo's NvExt that we care about.""" + model_config = {"extra": "ignore"} + fps: int | None = None + num_frames: int | None = None + seed: int | None = None + + +class Cosmos3VideoRequest(BaseModel): + model: str | None = None + prompt: str + # Direct fields (handler → worker pod, or standalone mode) + num_frames: int = 189 + resolution: str = "720p" + aspect_ratio: str | None = None + fps: int = 24 + seed: int | None = None + image_url: str | None = None + video_url: str | None = None + # Dynamo frontend passthrough fields (NvCreateVideoRequest) + size: str | None = None + input_reference: str | None = None + nvext: NvExtFields | None = None + # Action modality fields + model_mode: str | None = None + action_url: str | None = None + domain_name: str | None = None + action_chunk_size: int | None = None + raw_action_dim: int | None = None + image_size: int | None = None + + def resolve(self) -> tuple[int, str, str | None, int, int | None, str | None]: + """Resolve effective (num_frames, resolution, aspect_ratio, fps, seed, vision_url). + + Dynamo frontend fields (nvext, size, input_reference) take priority when + present, falling back to direct fields for standalone/backward-compat. + """ + nf = self.num_frames + res = self.resolution + ar = self.aspect_ratio + fps = self.fps + seed = self.seed + if self.nvext is not None: + if self.nvext.num_frames is not None: + nf = self.nvext.num_frames + if self.nvext.fps is not None: + fps = self.nvext.fps + if self.nvext.seed is not None: + seed = self.nvext.seed + if self.size is not None: + parts = self.size.lower().split("x") + if len(parts) == 2: + try: + w, h = int(parts[0]), int(parts[1]) + if (w, h) in _SIZE_TO_RES_AR: + res, ar = _SIZE_TO_RES_AR[(w, h)] + except ValueError: + pass + vision_url = self.video_url or self.image_url or self.input_reference + return nf, res, ar, fps, seed, vision_url + + +class Cosmos3VideoDataItem(BaseModel): + output_format: str = "mp4" + b64_json: str | None = None + url: str | None = None + seed: int | None = None + media_type: str = "video" + action: list | None = None + + +class Cosmos3VideoResponse(BaseModel): + """Matches Dynamo's NvVideosResponse so the frontend can parse it.""" + id: str = Field(default_factory=lambda: f"video-{uuid.uuid4().hex}") + object: str = "video" + model: str = "" + status: str = "completed" + progress: int = 100 + created: int = Field(default_factory=lambda: int(time.time())) + data: list[Cosmos3VideoDataItem] = Field(default_factory=list) + error: str | None = None + inference_time_s: float | None = None + + +# --------------------------------------------------------------------------- +# Inference engine (shared between both modes) +# --------------------------------------------------------------------------- + +_pipe: OmniInference | None = None +_served_model_name: str = "" +_output_root: Path = Path("/tmp/cosmos3_worker_outputs") +_generate_lock = threading.Lock() + + +def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: + """Run one generation collectively on all ranks. + + Returns the response dict on rank 0; ``None`` on other ranks. + """ + assert _pipe is not None + rank_dir = _output_root / payload["job_id"] / f"rank{get_rank()}" + rank_dir.mkdir(parents=True, exist_ok=True) + try: + vision_url = payload["vision_url"] + if vision_url and vision_url.startswith("data:"): + vision_url = _decode_data_uri(vision_url, rank_dir / "inputs") + + override_kwargs: dict[str, Any] = { + "name": "sample", + "output_dir": rank_dir, + "prompt": payload["prompt"], + "num_frames": payload["num_frames"], + "resolution": payload["resolution"], + "fps": payload["fps"], + "seed": payload["seed"], + "vision_path": vision_url, + } + if payload.get("aspect_ratio") is not None: + override_kwargs["aspect_ratio"] = payload["aspect_ratio"] + # Action modality fields (only set when present). + model_mode = payload.get("model_mode") + if model_mode is not None: + override_kwargs["model_mode"] = model_mode + if payload.get("action_url") is not None: + override_kwargs["action_path"] = payload["action_url"] + for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size"): + if payload.get(key) is not None: + override_kwargs[key] = payload[key] + + overrides = OmniSampleOverrides(**override_kwargs) + overrides.download(rank_dir / "inputs") + sample_args = overrides.build_sample(model_config=_pipe.model_config) + results = _pipe.generate([sample_args]) + + if not is_rank0(): + return None + if not results: + raise RuntimeError("No output produced") + result = results[0] + if result.status != "success": + raise RuntimeError(result.message or f"Generation {result.status}") + + # Extract action output if present (inverse_dynamics, policy modes). + action_output: list | None = None + for out in result.outputs: + if "action" in out.content: + action_output = out.content["action"] + break + + files = [f for out in result.outputs for f in out.files] + if not files: + raise RuntimeError("Output produced no media file") + output_file = Path(files[0]) + if not output_file.is_absolute(): + output_file = rank_dir / output_file + + b64 = base64.b64encode(output_file.read_bytes()).decode("ascii") + media_type = "image" if output_file.suffix.lower() in (".jpg", ".jpeg", ".png") else "video" + resp: dict[str, Any] = {"b64_json": b64, "seed": payload["seed"], "media_type": media_type} + if action_output is not None: + resp["action"] = action_output + return resp + finally: + shutil.rmtree(rank_dir, ignore_errors=True) + + +def _generate_blocking(request: Cosmos3VideoRequest) -> Cosmos3VideoResponse: + """Thread-safe synchronous generation. Called from async context via to_thread.""" + num_frames, resolution, aspect_ratio, fps, seed_val, vision_url = request.resolve() + seed = seed_val if seed_val is not None else random.randint(0, 2**31 - 1) + payload: dict[str, Any] = { + "job_id": uuid.uuid4().hex, + "prompt": request.prompt, + "num_frames": num_frames, + "resolution": _normalize_resolution(resolution), + "fps": fps, + "seed": seed, + "vision_url": vision_url, + } + if aspect_ratio is not None: + payload["aspect_ratio"] = aspect_ratio + # Forward action fields when present. + if request.model_mode is not None: + payload["model_mode"] = request.model_mode + if request.action_url is not None: + payload["action_url"] = request.action_url + for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size"): + val = getattr(request, key) + if val is not None: + payload[key] = val + + with _generate_lock: + distributed.broadcast_object(payload, src=0, group=_CONTROL_PG) + result = _run_generate(payload) + if result is None: + raise RuntimeError("rank 0 got None result") + media_type = result.get("media_type", "video") + output_format = "jpeg" if media_type == "image" else "mp4" + return Cosmos3VideoResponse( + model=_served_model_name, + data=[Cosmos3VideoDataItem( + output_format=output_format, + b64_json=result["b64_json"], + seed=result["seed"], + media_type=media_type, + action=result.get("action"), + )], + ) + + +# --------------------------------------------------------------------------- +# Multi-GPU: non-rank-0 worker loop +# --------------------------------------------------------------------------- + + +def _worker_rank_loop() -> None: + """Non-zero ranks: wait for broadcasts from rank 0 and join each generation.""" + torch.cuda.set_device(get_local_rank()) + while True: + # Blocks until rank 0 dispatches the next request. Runs on the gloo + # control group: an idle wait here may legitimately last hours. + payload = distributed.broadcast_object(None, src=0, group=_CONTROL_PG) + if payload is None: + continue + try: + _run_generate(payload) + except Exception as e: # noqa: BLE001 + log.error(f"rank {get_rank()} generation error: {e}") + + +# --------------------------------------------------------------------------- +# Mode A: Dynamo backend (production) +# --------------------------------------------------------------------------- + + +def _get_worker_namespace() -> str: + """Resolve Dynamo namespace for endpoint registration.""" + namespace = os.environ.get("DYN_NAMESPACE", "dynamo") + suffix = os.environ.get("DYN_NAMESPACE_WORKER_SUFFIX") + if suffix: + namespace = f"{namespace}-{suffix}" + return namespace + + +async def _register_model(endpoint, served_name: str, model_path: str) -> None: + from dynamo.llm import ModelInput, ModelType, register_llm # type: ignore[attr-defined] + + try: + await register_llm( + ModelInput.Text, + ModelType.Videos, + endpoint, + model_path, + served_name, + ) + log.info(f"Registered model with Dynamo: {served_name} (path={model_path})") + except Exception as e: + log.error(f"Failed to register model with Dynamo: {e}") + raise RuntimeError("Model registration failed") from e + + +async def _dynamo_main(args: argparse.Namespace) -> None: + """Dynamo backend main loop (rank 0 only).""" + from dynamo.runtime import DistributedRuntime, dynamo_endpoint + + loop = asyncio.get_running_loop() + discovery_backend = os.environ.get("DYN_DISCOVERY_BACKEND") + if not discovery_backend: + discovery_backend = "kubernetes" if os.environ.get("KUBERNETES_SERVICE_HOST") else "file" + namespace_name = _get_worker_namespace() + log.info(f"Dynamo discovery backend: {discovery_backend}") + log.info(f"Dynamo namespace: {namespace_name}") + + # DistributedRuntime(loop, discovery_backend, request_plane, enable_nats) + # enable_nats=False: DeepInfra clusters use ZMQ, not NATS. + runtime = DistributedRuntime(loop, discovery_backend, "tcp", False) + + component_name = "backend" + endpoint_name = "generate" + endpoint = runtime.endpoint(f"{namespace_name}.{component_name}.{endpoint_name}") + log.info(f"Serving Dynamo endpoint: {namespace_name}/{component_name}/{endpoint_name}") + + # Decorate the generate function with dynamo_endpoint for serialization. + @dynamo_endpoint(Cosmos3VideoRequest, Cosmos3VideoResponse) + async def create_video(request: Cosmos3VideoRequest): + if request.model is not None and request.model != _served_model_name: + raise ValueError(f"Model {request.model!r} not found. This worker serves {_served_model_name!r}.") + result = await asyncio.to_thread(_generate_blocking, request) + yield result.model_dump() + + # Log system port status for observability. + _system_port = os.environ.get("DYN_SYSTEM_PORT", "-1") + try: + _system_port_int = int(_system_port) + except ValueError: + _system_port_int = -1 + if _system_port_int < 0: + log.warning(f"DYN_SYSTEM_PORT not set (or is {_system_port!r}); metrics will NOT be exposed on a system port.") + else: + log.info(f"Dynamo system port enabled: {_system_port_int}") + + model_path = args.model + await asyncio.gather( + endpoint.serve_endpoint(create_video), + _register_model(endpoint, _served_model_name, model_path), + ) + + +# --------------------------------------------------------------------------- +# Mode B: Standalone HTTP server (local testing) +# --------------------------------------------------------------------------- + + +def _run_standalone_http(args: argparse.Namespace) -> None: + """Run a plain FastAPI HTTP server (no Dynamo frontend needed).""" + import queue + + import uvicorn + from fastapi import FastAPI, HTTPException + + request_queue: "queue.Queue[dict[str, Any]]" = queue.Queue() + app = FastAPI(title="Cosmos 3 inference worker (standalone)") + + @app.get("/health") + def health() -> dict[str, str]: + return {"status": "ok"} + + @app.post("/v1/videos") + def create_video(req: Cosmos3VideoRequest) -> dict: + if req.model is not None and req.model != _served_model_name: + raise HTTPException( + status_code=404, + detail=f"Model {req.model!r} not found. This worker serves {_served_model_name!r}.", + ) + num_frames, resolution, aspect_ratio, fps, seed_val, vision_url = req.resolve() + seed = seed_val if seed_val is not None else random.randint(0, 2**31 - 1) + payload = { + "job_id": uuid.uuid4().hex, + "prompt": req.prompt, + "num_frames": num_frames, + "resolution": _normalize_resolution(resolution), + "fps": fps, + "seed": seed, + "vision_url": vision_url, + } + if aspect_ratio is not None: + payload["aspect_ratio"] = aspect_ratio + # Forward action fields when present. + if req.model_mode is not None: + payload["model_mode"] = req.model_mode + if req.action_url is not None: + payload["action_url"] = req.action_url + for key in ("domain_name", "action_chunk_size", "raw_action_dim", "image_size"): + val = getattr(req, key) + if val is not None: + payload[key] = val + job: dict[str, Any] = { + "payload": payload, + "event": threading.Event(), + "result": None, + "error": None, + } + request_queue.put(job) + job["event"].wait() + + if job["error"] is not None: + raise HTTPException(status_code=400, detail=str(job["error"])) + return job["result"] + + def inference_thread() -> None: + torch.cuda.set_device(get_local_rank()) + while True: + job = request_queue.get() + try: + distributed.broadcast_object(job["payload"], src=0, group=_CONTROL_PG) + result = _run_generate(job["payload"]) + media_type = result.get("media_type", "video") if result else "video" + output_format = "jpeg" if media_type == "image" else "mp4" + item = { + "output_format": output_format, + "b64_json": result["b64_json"], + "seed": result["seed"], + "media_type": media_type, + } + if result.get("action") is not None: + item["action"] = result["action"] + job["result"] = Cosmos3VideoResponse( + model=_served_model_name, + data=[Cosmos3VideoDataItem(**item)], + ).model_dump() + except Exception as e: # noqa: BLE001 + log.error(f"Generation failed: {e}") + job["error"] = e + finally: + job["event"].set() + + threading.Thread(target=inference_thread, daemon=True).start() + log.success(f"Starting standalone HTTP server on {args.host}:{args.port}") + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +# --------------------------------------------------------------------------- +# CLI and main +# --------------------------------------------------------------------------- + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Cosmos 3 Dynamo videogen worker") + parser.add_argument( + "--served-model-name", + required=True, + help="Public model name clients request, e.g. nvidia/Cosmos3-Nano", + ) + parser.add_argument( + "--model", + required=True, + help="Checkpoint path or registered name, e.g. Cosmos3-Nano or /data/default", + ) + parser.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs for this worker") + parser.add_argument( + "--standalone", + action="store_true", + help="Run as standalone HTTP server (no Dynamo frontend required)", + ) + parser.add_argument("--host", default="0.0.0.0", help="Host (standalone mode only)") + parser.add_argument("--port", type=int, default=80, help="Port (standalone mode only)") + return parser.parse_args() + + +def _relaunch_under_torchrun(num_gpus: int) -> None: + """Re-exec under torchrun for multi-GPU collective inference.""" + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={num_gpus}", + os.path.abspath(__file__), + *sys.argv[1:], + ] + log.info(f"Relaunching under torchrun for {num_gpus} GPUs: {' '.join(cmd)}") + os.execvp(cmd[0], cmd) + + +def main() -> None: + global _pipe, _served_model_name, _output_root + + args = parse_args() + + # Multi-GPU: relaunch under torchrun unless already a torchrun child. + if args.num_gpus > 1 and "RANK" not in os.environ: + _relaunch_under_torchrun(args.num_gpus) + + _served_model_name = args.served_model_name + _output_root = Path(os.environ.get("COSMOS3_OUTPUT_DIR", "/tmp/cosmos3_worker_outputs")).absolute() + init_output_dir(_output_root) + + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + log.info( + f"Loading Cosmos 3 model {_served_model_name!r} from checkpoint {args.model!r} (world_size={world_size})..." + ) + setup_args = OmniSetupOverrides( + checkpoint_path=args.model, + output_dir=_output_root, + ).build_setup(world_size=world_size) + _pipe = OmniInference.create(setup_args) + log.success("Model loaded.") + + if not is_rank0(): + _worker_rank_loop() + return + + # Rank 0: serve requests. + if args.standalone: + _run_standalone_http(args) + else: + import uvloop + + log.success("Starting Dynamo backend...") + uvloop.install() + asyncio.run(_dynamo_main(args)) + + +if __name__ == "__main__": + main() From 8fabfc06d6adcf7910e70535373d0b640d7d55f2 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Tue, 7 Jul 2026 23:14:41 +0000 Subject: [PATCH 2/5] cosmos3: pass HF token via BuildKit secret mount, not build ARG ARG values are recorded in image layer history, so a token passed that way is readable by anyone who can pull from the registry. The secret mount is available only during the RUN and never persists. Co-Authored-By: Claude Fable 5 --- examples/cosmos3/Dockerfile | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/examples/cosmos3/Dockerfile b/examples/cosmos3/Dockerfile index 706778f6bb8a..7b2f136c9156 100644 --- a/examples/cosmos3/Dockerfile +++ b/examples/cosmos3/Dockerfile @@ -6,8 +6,10 @@ # There is no DeepInfra copy of the framework repo — this directory plus the # pin below is the complete, reviewable definition of the image. # -# Build (from this directory): -# docker build --build-arg HF_TOKEN= \ +# Build (from this directory; token file must hold an HF token whose account +# has accepted the nvidia/Cosmos-Guardrail1 auto-gate — the secret mount keeps +# it out of image layers and history): +# docker build --secret id=hf_token,src=/path/to/token-file \ # -t localhost:30500/cosmos3-worker:v1 . ARG CUDA_VERSION=13.0.2 @@ -73,46 +75,48 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Pre-fetch HuggingFace dependencies needed by cosmos_framework at model init # time: Qwen tokenizer configs (*.json, *.txt — few MB), Wan VAE checkpoints, -# and the full Cosmos-Guardrail1 snapshot (gated repo: the build HF_TOKEN's -# account must have accepted access; gating is auto-approve). +# and the full Cosmos-Guardrail1 snapshot (the one gated repo — auto-approve +# gate; the token account must have accepted it). The token comes in via a +# BuildKit secret mount so it never persists in a layer or in image history. # This eliminates any runtime dependency on HuggingFace availability. # Revisions MUST match the pins in # cosmos_framework/inference/common/checkpoints.py and # cosmos_framework/auxiliary/guardrail/common/core.py at the ref above — the # runtime looks up these exact revisions; a mismatch silently falls back to a # live (unauthenticated, rate-limited) HF fetch at every pod boot. -ARG HF_TOKEN="" -RUN HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ +RUN --mount=type=secret,id=hf_token \ + export HF_TOKEN=$(cat /run/secrets/hf_token 2>/dev/null || true) && \ + uvx hf@1.13.0 download --format=json \ Qwen/Qwen3-0.6B \ --repo-type model \ --revision c1899de289a04d12100db370d81485cdf75e47ca \ --include '*.json' --include '*.txt' && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ Qwen/Qwen3-VL-2B-Instruct \ --repo-type model \ --revision 89644892e4d85e24eaac8bacfd4f463576704203 \ --include '*.json' --include '*.txt' && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ Qwen/Qwen3-VL-8B-Instruct \ --repo-type model \ --revision 0c351dd01ed87e9c1b53cbc748cba10e6187ff3b \ --include '*.json' --include '*.txt' && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ Qwen/Qwen3-VL-32B-Instruct \ --repo-type model \ --revision 0cfaf48183f594c314753d30a4c4974bc75f3ccb \ --include '*.json' --include '*.txt' && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ Wan-AI/Wan2.1-T2V-14B \ --repo-type model \ --revision a064a6c71f5be440641209c07bf2a5ce7a2ff5e4 \ Wan2.1_VAE.pth && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ Wan-AI/Wan2.2-TI2V-5B \ --repo-type model \ --revision 921dbaf3f1674a56f47e83fb80a34bac8a8f203e \ Wan2.2_VAE.pth && \ - HF_TOKEN=${HF_TOKEN} uvx hf@1.13.0 download --format=json \ + uvx hf@1.13.0 download --format=json \ nvidia/Cosmos-Guardrail1 \ --repo-type model \ --revision d6d4bfa899a71454a700907664f3e88f503950cf From 06ffe166d7f1915495103bce3d4fc5311c641411 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Wed, 8 Jul 2026 04:24:23 +0000 Subject: [PATCH 3/5] cosmos3: disable face blur, add BT.709 signaling, reorder patches last MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - patches/0002: video guardrail runner drops the RetinaFace face-blur postprocessor, which unconditionally pixelates every detected face in every output frame (breaks the product; it is a privacy/anti-deepfake measure, not a content-safety control). The content-safety value stays: text-prompt guardrail (Blocklist + Qwen3Guard) is unchanged, and upstream already disables VideoContentSafetyFilter. Removal explicitly authorized by Johan (product owner) 2026-07-08. Flagged as potentially relevant to likeness/NCII/deepfake obligations; to be confirmed with legal as a follow-up (not a rollout blocker per Johan). - patches/0001 (encode-params): -threads cap plus explicit BT.709 signaling (colorspace+primaries+trc + limited range); conversion target and tag set together. Verified to round-trip identically to untagged on a compliant decoder — correctness hygiene, not a fix for a confirmed visible defect. - Dockerfile: apply patches + copy worker LAST (after the ~11GB gated prefetch), relying on the editable install, so patch iterations rebuild in seconds. Co-Authored-By: Claude Fable 5 --- examples/cosmos3/Dockerfile | 15 +++++-- .../patches/0001-cap-encoder-threads.patch | 33 -------------- .../cosmos3/patches/0001-encode-params.patch | 43 +++++++++++++++++++ .../patches/0002-disable-face-blur.patch | 29 +++++++++++++ 4 files changed, 83 insertions(+), 37 deletions(-) delete mode 100644 examples/cosmos3/patches/0001-cap-encoder-threads.patch create mode 100644 examples/cosmos3/patches/0001-encode-params.patch create mode 100644 examples/cosmos3/patches/0002-disable-face-blur.patch diff --git a/examples/cosmos3/Dockerfile b/examples/cosmos3/Dockerfile index 7b2f136c9156..944062817b66 100644 --- a/examples/cosmos3/Dockerfile +++ b/examples/cosmos3/Dockerfile @@ -44,12 +44,12 @@ WORKDIR /workspace RUN git clone --filter=blob:none ${COSMOS_FRAMEWORK_REPO} /workspace && \ git -C /workspace checkout ${COSMOS_FRAMEWORK_REF} -# DeepInfra patches to upstream code (reviewed in deepinfra/dynamo). -COPY patches /workspace/deepinfra-patches -RUN git -C /workspace apply --verbose deepinfra-patches/*.patch - # Upstream build steps (mirrors NVIDIA/cosmos-framework Dockerfile @ the pin, # minus bind-mounts: the sources are already in the layer from the clone). +# NOTE: patches are applied LAST (near the entrypoint), not here. cosmos_framework +# is installed editable (`uv pip install -e`), so patching /workspace afterward is +# picked up live at runtime — and keeping the patch layer last means editing a +# patch rebuilds in seconds instead of re-running the ~11GB gated HF prefetch. RUN --mount=type=cache,target=/root/.cache/uv uv python install RUN echo "$CUDA_VERSION" | sed -E 's/^([0-9]+)\.([0-9]+).*/cu\1\2/' > /root/.cuda-name RUN --mount=type=cache,target=/root/.cache/uv \ @@ -130,6 +130,13 @@ ENV TORCHINDUCTOR_CACHE_DIR="/cache/torchinductor" ENV TRITON_CACHE_DIR="/cache/triton" RUN mkdir -p /cache/torchinductor /cache/triton +# DeepInfra patches to upstream code (reviewed in deepinfra/dynamo). Applied +# last so a patch edit doesn't invalidate the expensive prefetch layer above; +# the editable install makes them live at runtime. Fails loudly if a patch no +# longer applies against COSMOS_FRAMEWORK_REF (the signal to rebase or drop it). +COPY patches /workspace/deepinfra-patches +RUN git -C /workspace apply --verbose deepinfra-patches/*.patch + # DeepInfra Dynamo videogen worker. Runtime args # (--served-model-name, --model, --num-gpus, ...) are appended to this command. # worker.py self-relaunches under torchrun when --num-gpus > 1. diff --git a/examples/cosmos3/patches/0001-cap-encoder-threads.patch b/examples/cosmos3/patches/0001-cap-encoder-threads.patch deleted file mode 100644 index 3cd17fcc7b58..000000000000 --- a/examples/cosmos3/patches/0001-cap-encoder-threads.patch +++ /dev/null @@ -1,33 +0,0 @@ -diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -index fdc944f..8393fa5 100644 ---- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -+++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -@@ -1,6 +1,7 @@ - # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - # SPDX-License-Identifier: OpenMDW-1.1 - -+import os - from typing import IO, Any, Dict, Tuple - - import imageio -@@ -171,6 +172,20 @@ class ImageioVideoHandler(BaseFileHandler): - } - # Update with any other kwargs - mimsave_kwargs.update(kwargs) -+ -+ # deepinfra: cap CPU encoder threads. imageio encodes mp4 with CPU -+ # libx264; with no -threads, x264 auto-detects every core and -+ # oversubscribes on shared multi-GPU nodes (e.g. 8 GPUs / 224 cores), -+ # contending with neighbouring GPU pods so encode latency drifts with -+ # node load. Append a per-GPU-sized cap unless a caller already -+ # specified -threads. Override via DI_VIDEO_ENCODE_THREADS. -+ if "-threads" not in mimsave_kwargs.get("ffmpeg_params", []): -+ mimsave_kwargs["ffmpeg_params"] = [ -+ *mimsave_kwargs.get("ffmpeg_params", []), -+ "-threads", -+ os.getenv("DI_VIDEO_ENCODE_THREADS", "32"), -+ ] -+ - log.debug(f"mimsave_kwargs: {mimsave_kwargs}") - - imageio.mimsave(file, obj, format, **mimsave_kwargs) diff --git a/examples/cosmos3/patches/0001-encode-params.patch b/examples/cosmos3/patches/0001-encode-params.patch new file mode 100644 index 000000000000..6f37dcad9cd0 --- /dev/null +++ b/examples/cosmos3/patches/0001-encode-params.patch @@ -0,0 +1,43 @@ +diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +index fdc944f..a3f373e 100644 +--- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py ++++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +@@ -1,6 +1,7 @@ + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + # SPDX-License-Identifier: OpenMDW-1.1 + ++import os + from typing import IO, Any, Dict, Tuple + + import imageio +@@ -171,6 +172,30 @@ class ImageioVideoHandler(BaseFileHandler): + } + # Update with any other kwargs + mimsave_kwargs.update(kwargs) ++ ++ # deepinfra: adjust the libx264 encode params (covers both the crf and ++ # qscale branches above, plus any caller-supplied ffmpeg_params). ++ _params = list(mimsave_kwargs.get("ffmpeg_params") or []) ++ # (1) Cap CPU encoder threads. With no -threads, x264 auto-detects every ++ # core and oversubscribes on shared multi-GPU nodes (8 GPUs / 224 ++ # cores), contending with neighbouring pods so encode latency drifts ++ # with node load. Override via DI_VIDEO_ENCODE_THREADS. ++ if "-threads" not in _params: ++ _params += ["-threads", os.getenv("DI_VIDEO_ENCODE_THREADS", "32")] ++ # (2) Signal BT.709 explicitly (matrix + transfer + primaries + limited ++ # range) instead of leaving the stream untagged and relying on the ++ # player's resolution-based guess. The conversion target and the tag ++ # are set together so they cannot disagree. Correctness hygiene, not ++ # a fix for any confirmed visible defect. ++ if "-colorspace" not in _params: ++ _params += [ ++ "-colorspace", "bt709", ++ "-color_primaries", "bt709", ++ "-color_trc", "bt709", ++ "-color_range", "tv", ++ ] ++ mimsave_kwargs["ffmpeg_params"] = _params ++ + log.debug(f"mimsave_kwargs: {mimsave_kwargs}") + + imageio.mimsave(file, obj, format, **mimsave_kwargs) diff --git a/examples/cosmos3/patches/0002-disable-face-blur.patch b/examples/cosmos3/patches/0002-disable-face-blur.patch new file mode 100644 index 000000000000..c22fda0a9dd5 --- /dev/null +++ b/examples/cosmos3/patches/0002-disable-face-blur.patch @@ -0,0 +1,29 @@ +diff --git a/cosmos_framework/auxiliary/guardrail/common/presets.py b/cosmos_framework/auxiliary/guardrail/common/presets.py +index 8536685..932daa5 100644 +--- a/cosmos_framework/auxiliary/guardrail/common/presets.py ++++ b/cosmos_framework/auxiliary/guardrail/common/presets.py +@@ -21,13 +21,22 @@ def create_text_guardrail_runner(offload_model_to_cpu: bool = False) -> Guardrai + + + def create_video_guardrail_runner(offload_model_to_cpu: bool = False) -> GuardrailRunner: +- """Create the video guardrail runner.""" ++ """Create the video guardrail runner. ++ ++ deepinfra: the RetinaFace face-blur postprocessor is intentionally removed. ++ It unconditionally pixelates every detected face in every output frame, ++ which breaks the product (people are a core use case) and is a privacy/ ++ anti-deepfake measure, not a content-safety control. The content-safety ++ value lives in the text-prompt guardrail (Blocklist + Qwen3Guard), which is ++ unchanged. Upstream's VideoContentSafetyFilter is already disabled by NVIDIA ++ (false positives), so the video runner is now a no-op passthrough. ++ """ + return GuardrailRunner( + safety_models=[ + # VideoContentSafetyFilter(offload_model_to_cpu=offload_model_to_cpu) + # Too many false positives, add back when fixed + ], +- postprocessors=[RetinaFaceFilter(offload_model_to_cpu=offload_model_to_cpu)], ++ postprocessors=[], + ) + + From f12acdf599a0340b73aafc50561a6e4dedda6f86 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Wed, 8 Jul 2026 21:07:08 +0000 Subject: [PATCH 4/5] =?UTF-8?q?cosmos3:=20drop=20the=20BT.709=20color=20pa?= =?UTF-8?q?tch=20=E2=80=94=20it=20degraded=20color,=20not=20hygiene?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In-image A/B on a color test card, both decoded as BT.709 by a compliant decoder: fw-v1 untagged (current): red 253,0,0 green 0,255,0 skin 239,193,165 (~exact) bt709/tv tagged (patch): red 255,23,0 green 0,215,0 skin 243,196,163 (err up to 40) Setting -colorspace/-color_range as imageio output flags TAGS the stream but does not correctly drive the RGB->YUV conversion, so tag and pixels disagree (the mislabel trap). The untagged output already round-trips cleanly on a compliant decoder, so there is no confirmed visible defect to fix here anyway. 0001 reverts to the thread-cap-only change. Proper color signaling (force the conversion via a scale filter matched to the tag) is deferred to its own task, gated on first confirming an actual browser-rendering problem exists. fw-v2 = face-blur removal (0002) + encoder thread cap (0001). No color change. Co-Authored-By: Claude Fable 5 --- .../patches/0001-cap-encoder-threads.patch | 30 +++++++++++++ .../cosmos3/patches/0001-encode-params.patch | 43 ------------------- 2 files changed, 30 insertions(+), 43 deletions(-) create mode 100644 examples/cosmos3/patches/0001-cap-encoder-threads.patch delete mode 100644 examples/cosmos3/patches/0001-encode-params.patch diff --git a/examples/cosmos3/patches/0001-cap-encoder-threads.patch b/examples/cosmos3/patches/0001-cap-encoder-threads.patch new file mode 100644 index 000000000000..e3e2cf2c24e7 --- /dev/null +++ b/examples/cosmos3/patches/0001-cap-encoder-threads.patch @@ -0,0 +1,30 @@ +diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +index fdc944f..359d016 100644 +--- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py ++++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py +@@ -1,6 +1,7 @@ + # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + # SPDX-License-Identifier: OpenMDW-1.1 + ++import os + from typing import IO, Any, Dict, Tuple + + import imageio +@@ -171,6 +172,17 @@ class ImageioVideoHandler(BaseFileHandler): + } + # Update with any other kwargs + mimsave_kwargs.update(kwargs) ++ ++ # deepinfra: cap CPU encoder threads (covers the crf and qscale branches ++ # above, plus any caller-supplied ffmpeg_params). With no -threads, x264 ++ # auto-detects every core and oversubscribes on shared multi-GPU nodes ++ # (8 GPUs / 224 cores), contending with neighbouring pods so encode ++ # latency drifts with node load. Override via DI_VIDEO_ENCODE_THREADS. ++ _params = list(mimsave_kwargs.get("ffmpeg_params") or []) ++ if "-threads" not in _params: ++ _params += ["-threads", os.getenv("DI_VIDEO_ENCODE_THREADS", "32")] ++ mimsave_kwargs["ffmpeg_params"] = _params ++ + log.debug(f"mimsave_kwargs: {mimsave_kwargs}") + + imageio.mimsave(file, obj, format, **mimsave_kwargs) diff --git a/examples/cosmos3/patches/0001-encode-params.patch b/examples/cosmos3/patches/0001-encode-params.patch deleted file mode 100644 index 6f37dcad9cd0..000000000000 --- a/examples/cosmos3/patches/0001-encode-params.patch +++ /dev/null @@ -1,43 +0,0 @@ -diff --git a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -index fdc944f..a3f373e 100644 ---- a/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -+++ b/cosmos_framework/utils/easy_io/handlers/imageio_video_handler.py -@@ -1,6 +1,7 @@ - # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - # SPDX-License-Identifier: OpenMDW-1.1 - -+import os - from typing import IO, Any, Dict, Tuple - - import imageio -@@ -171,6 +172,30 @@ class ImageioVideoHandler(BaseFileHandler): - } - # Update with any other kwargs - mimsave_kwargs.update(kwargs) -+ -+ # deepinfra: adjust the libx264 encode params (covers both the crf and -+ # qscale branches above, plus any caller-supplied ffmpeg_params). -+ _params = list(mimsave_kwargs.get("ffmpeg_params") or []) -+ # (1) Cap CPU encoder threads. With no -threads, x264 auto-detects every -+ # core and oversubscribes on shared multi-GPU nodes (8 GPUs / 224 -+ # cores), contending with neighbouring pods so encode latency drifts -+ # with node load. Override via DI_VIDEO_ENCODE_THREADS. -+ if "-threads" not in _params: -+ _params += ["-threads", os.getenv("DI_VIDEO_ENCODE_THREADS", "32")] -+ # (2) Signal BT.709 explicitly (matrix + transfer + primaries + limited -+ # range) instead of leaving the stream untagged and relying on the -+ # player's resolution-based guess. The conversion target and the tag -+ # are set together so they cannot disagree. Correctness hygiene, not -+ # a fix for any confirmed visible defect. -+ if "-colorspace" not in _params: -+ _params += [ -+ "-colorspace", "bt709", -+ "-color_primaries", "bt709", -+ "-color_trc", "bt709", -+ "-color_range", "tv", -+ ] -+ mimsave_kwargs["ffmpeg_params"] = _params -+ - log.debug(f"mimsave_kwargs: {mimsave_kwargs}") - - imageio.mimsave(file, obj, format, **mimsave_kwargs) From 78aa7cf77e386aff89755b4380dcf9b88e81317c Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Thu, 9 Jul 2026 20:14:26 +0000 Subject: [PATCH 5/5] cosmos3: satisfy pre-commit (black/isort format + noqa: E402 on deferred imports) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Formatting-only: black + isort reflow, and # noqa: E402 on the imports that must follow init_script() (the required init-before-heavy-imports pattern, same as cosmos_framework's own scripts). No behavior change — the built fw-v2 image (from the pre-format worker.py) is byte-behavior-identical. Co-Authored-By: Claude Fable 5 --- examples/cosmos3/worker.py | 123 +++++++++++++++++++++++++------------ 1 file changed, 83 insertions(+), 40 deletions(-) diff --git a/examples/cosmos3/worker.py b/examples/cosmos3/worker.py index daa76fcf0d6f..f341a28a24bd 100644 --- a/examples/cosmos3/worker.py +++ b/examples/cosmos3/worker.py @@ -59,28 +59,35 @@ init_script() -import argparse -import asyncio -import base64 -import mimetypes -import os -import random -import re -import shutil -import sys -import threading -import time -import uuid -from pathlib import Path -from typing import Any - -import torch -from pydantic import BaseModel, Field - -from cosmos_framework.inference.args import OmniSampleOverrides, OmniSetupOverrides -from cosmos_framework.inference.common.init import get_local_rank, get_rank, init_output_dir, is_rank0 -from cosmos_framework.inference.inference import OmniInference -from cosmos_framework.utils import distributed, log +import argparse # noqa: E402 +import asyncio # noqa: E402 +import base64 # noqa: E402 +import mimetypes # noqa: E402 +import os # noqa: E402 +import random # noqa: E402 +import re # noqa: E402 +import shutil # noqa: E402 +import sys # noqa: E402 +import threading # noqa: E402 +import time # noqa: E402 +import uuid # noqa: E402 +from pathlib import Path # noqa: E402 +from typing import Any # noqa: E402 + +import torch # noqa: E402 +from cosmos_framework.inference.args import ( # noqa: E402 + OmniSampleOverrides, + OmniSetupOverrides, +) +from cosmos_framework.inference.common.init import ( # noqa: E402 + get_local_rank, + get_rank, + init_output_dir, + is_rank0, +) +from cosmos_framework.inference.inference import OmniInference # noqa: E402 +from cosmos_framework.utils import distributed, log # noqa: E402 +from pydantic import BaseModel, Field # noqa: E402 # --------------------------------------------------------------------------- # Control-plane process group (multi-GPU request dispatch) @@ -101,7 +108,8 @@ import datetime as _dt _CONTROL_PG = torch.distributed.new_group( - backend="gloo", timeout=_dt.timedelta(days=7)) + backend="gloo", timeout=_dt.timedelta(days=7) + ) # --------------------------------------------------------------------------- # Constants @@ -115,11 +123,16 @@ # for resolutions we support). Used to decode the Dynamo frontend's "size" field # ("WxH") back to the resolution + aspect_ratio the inference engine expects. _SIZE_TO_RES_AR: dict[tuple[int, int], tuple[str, str]] = {} + + def _init_size_lookup() -> None: from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO + for res, ar_map in VIDEO_RES_SIZE_INFO.items(): for ar, (w, h) in ar_map.items(): _SIZE_TO_RES_AR[(w, h)] = (res, ar) + + _init_size_lookup() @@ -159,6 +172,7 @@ def _decode_data_uri(uri: str, output_dir: Path) -> str: class NvExtFields(BaseModel): """Subset of Dynamo's NvExt that we care about.""" + model_config = {"extra": "ignore"} fps: int | None = None num_frames: int | None = None @@ -230,6 +244,7 @@ class Cosmos3VideoDataItem(BaseModel): class Cosmos3VideoResponse(BaseModel): """Matches Dynamo's NvVideosResponse so the frontend can parse it.""" + id: str = Field(default_factory=lambda: f"video-{uuid.uuid4().hex}") object: str = "video" model: str = "" @@ -314,8 +329,16 @@ def _run_generate(payload: dict[str, Any]) -> dict[str, Any] | None: output_file = rank_dir / output_file b64 = base64.b64encode(output_file.read_bytes()).decode("ascii") - media_type = "image" if output_file.suffix.lower() in (".jpg", ".jpeg", ".png") else "video" - resp: dict[str, Any] = {"b64_json": b64, "seed": payload["seed"], "media_type": media_type} + media_type = ( + "image" + if output_file.suffix.lower() in (".jpg", ".jpeg", ".png") + else "video" + ) + resp: dict[str, Any] = { + "b64_json": b64, + "seed": payload["seed"], + "media_type": media_type, + } if action_output is not None: resp["action"] = action_output return resp @@ -357,13 +380,15 @@ def _generate_blocking(request: Cosmos3VideoRequest) -> Cosmos3VideoResponse: output_format = "jpeg" if media_type == "image" else "mp4" return Cosmos3VideoResponse( model=_served_model_name, - data=[Cosmos3VideoDataItem( - output_format=output_format, - b64_json=result["b64_json"], - seed=result["seed"], - media_type=media_type, - action=result.get("action"), - )], + data=[ + Cosmos3VideoDataItem( + output_format=output_format, + b64_json=result["b64_json"], + seed=result["seed"], + media_type=media_type, + action=result.get("action"), + ) + ], ) @@ -402,7 +427,11 @@ def _get_worker_namespace() -> str: async def _register_model(endpoint, served_name: str, model_path: str) -> None: - from dynamo.llm import ModelInput, ModelType, register_llm # type: ignore[attr-defined] + from dynamo.llm import ( # type: ignore[attr-defined] + ModelInput, + ModelType, + register_llm, + ) try: await register_llm( @@ -425,7 +454,9 @@ async def _dynamo_main(args: argparse.Namespace) -> None: loop = asyncio.get_running_loop() discovery_backend = os.environ.get("DYN_DISCOVERY_BACKEND") if not discovery_backend: - discovery_backend = "kubernetes" if os.environ.get("KUBERNETES_SERVICE_HOST") else "file" + discovery_backend = ( + "kubernetes" if os.environ.get("KUBERNETES_SERVICE_HOST") else "file" + ) namespace_name = _get_worker_namespace() log.info(f"Dynamo discovery backend: {discovery_backend}") log.info(f"Dynamo namespace: {namespace_name}") @@ -437,13 +468,17 @@ async def _dynamo_main(args: argparse.Namespace) -> None: component_name = "backend" endpoint_name = "generate" endpoint = runtime.endpoint(f"{namespace_name}.{component_name}.{endpoint_name}") - log.info(f"Serving Dynamo endpoint: {namespace_name}/{component_name}/{endpoint_name}") + log.info( + f"Serving Dynamo endpoint: {namespace_name}/{component_name}/{endpoint_name}" + ) # Decorate the generate function with dynamo_endpoint for serialization. @dynamo_endpoint(Cosmos3VideoRequest, Cosmos3VideoResponse) async def create_video(request: Cosmos3VideoRequest): if request.model is not None and request.model != _served_model_name: - raise ValueError(f"Model {request.model!r} not found. This worker serves {_served_model_name!r}.") + raise ValueError( + f"Model {request.model!r} not found. This worker serves {_served_model_name!r}." + ) result = await asyncio.to_thread(_generate_blocking, request) yield result.model_dump() @@ -454,7 +489,9 @@ async def create_video(request: Cosmos3VideoRequest): except ValueError: _system_port_int = -1 if _system_port_int < 0: - log.warning(f"DYN_SYSTEM_PORT not set (or is {_system_port!r}); metrics will NOT be exposed on a system port.") + log.warning( + f"DYN_SYSTEM_PORT not set (or is {_system_port!r}); metrics will NOT be exposed on a system port." + ) else: log.info(f"Dynamo system port enabled: {_system_port_int}") @@ -575,14 +612,18 @@ def parse_args() -> argparse.Namespace: required=True, help="Checkpoint path or registered name, e.g. Cosmos3-Nano or /data/default", ) - parser.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs for this worker") + parser.add_argument( + "--num-gpus", type=int, default=1, help="Number of GPUs for this worker" + ) parser.add_argument( "--standalone", action="store_true", help="Run as standalone HTTP server (no Dynamo frontend required)", ) parser.add_argument("--host", default="0.0.0.0", help="Host (standalone mode only)") - parser.add_argument("--port", type=int, default=80, help="Port (standalone mode only)") + parser.add_argument( + "--port", type=int, default=80, help="Port (standalone mode only)" + ) return parser.parse_args() @@ -611,7 +652,9 @@ def main() -> None: _relaunch_under_torchrun(args.num_gpus) _served_model_name = args.served_model_name - _output_root = Path(os.environ.get("COSMOS3_OUTPUT_DIR", "/tmp/cosmos3_worker_outputs")).absolute() + _output_root = Path( + os.environ.get("COSMOS3_OUTPUT_DIR", "/tmp/cosmos3_worker_outputs") + ).absolute() init_output_dir(_output_root) world_size = int(os.environ.get("WORLD_SIZE", "1"))