Merge pull request #14 from deepinfra/johan/cosmos3-framework-worker

feat: Cosmos3 videogen worker on public cosmos-framework (overlay)

Author: Johan-de-R

examples/cosmos3/Dockerfile

@@ -0,0 +1,144 @@
+# 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; 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
+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}
+
+# 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 \
+    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 (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.
+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' && \
+    uvx hf@1.13.0 download --format=json \
+        Qwen/Qwen3-VL-2B-Instruct \
+        --repo-type model \
+        --revision 89644892e4d85e24eaac8bacfd4f463576704203 \
+        --include '*.json' --include '*.txt' && \
+    uvx hf@1.13.0 download --format=json \
+        Qwen/Qwen3-VL-8B-Instruct \
+        --repo-type model \
+        --revision 0c351dd01ed87e9c1b53cbc748cba10e6187ff3b \
+        --include '*.json' --include '*.txt' && \
+    uvx hf@1.13.0 download --format=json \
+        Qwen/Qwen3-VL-32B-Instruct \
+        --repo-type model \
+        --revision 0cfaf48183f594c314753d30a4c4974bc75f3ccb \
+        --include '*.json' --include '*.txt' && \
+    uvx hf@1.13.0 download --format=json \
+        Wan-AI/Wan2.1-T2V-14B \
+        --repo-type model \
+        --revision a064a6c71f5be440641209c07bf2a5ce7a2ff5e4 \
+        Wan2.1_VAE.pth && \
+    uvx hf@1.13.0 download --format=json \
+        Wan-AI/Wan2.2-TI2V-5B \
+        --repo-type model \
+        --revision 921dbaf3f1674a56f47e83fb80a34bac8a8f203e \
+        Wan2.2_VAE.pth && \
+    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 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.
+COPY worker.py /workspace/worker.py
+ENTRYPOINT ["python3", "worker.py"]

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.

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)

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=[],
+     )
+ 
+