From 3b0dbf22aafc89fbf8fe0e55f3528bd09a87c712 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Tue, 27 Oct 2026 10:49:06 -0700 Subject: [PATCH 1/2] Add Cosmos3 Diffusers CPU-GPU offloading package --- README.md | 2 +- packages/diffusers-cpu-offloading/.gitignore | 11 + packages/diffusers-cpu-offloading/README.md | 193 ++++ .../benchmarks/README.md | 210 ++++ .../benchmarks/benchmark_bf16.py | 6 + .../benchmarks/benchmark_fp8.py | 6 + .../benchmarks/benchmark_safety.py | 209 ++++ .../benchmarks/fp8_benchmark_audit.py | 96 ++ .../benchmarks/manifest.json | 94 ++ .../benchmarks/modalities.py | 198 ++++ .../benchmarks/runner.py | 630 ++++++++++++ .../benchmarks/v2v_negative_prompt.json | 108 ++ .../benchmarks/v2v_prompt.json | 136 +++ .../cosmos3_streaming/__init__.py | 26 + .../cosmos3_streaming/attention/__init__.py | 9 + .../cosmos3_streaming/attention/base.py | 24 + .../cosmos3_streaming/attention/loader.py | 22 + .../cosmos3_streaming/attention/sdpa.py | 53 + .../cosmos3_streaming/block_views.py | 335 ++++++ .../cosmos3_streaming/config.py | 43 + .../cosmos3_streaming/fp8.py | 121 +++ .../cosmos3_streaming/memory_trace.py | 93 ++ .../cosmos3_streaming/metrics.py | 22 + .../cosmos3_streaming/modular.py | 656 ++++++++++++ .../cosmos3_streaming/offloading.py | 51 + .../cosmos3_streaming/runtime.py | 970 ++++++++++++++++++ .../diffusers-cpu-offloading/pyproject.toml | 39 + 27 files changed, 4362 insertions(+), 1 deletion(-) create mode 100644 packages/diffusers-cpu-offloading/.gitignore create mode 100644 packages/diffusers-cpu-offloading/README.md create mode 100644 packages/diffusers-cpu-offloading/benchmarks/README.md create mode 100644 packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/manifest.json create mode 100644 packages/diffusers-cpu-offloading/benchmarks/modalities.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/runner.py create mode 100644 packages/diffusers-cpu-offloading/benchmarks/v2v_negative_prompt.json create mode 100644 packages/diffusers-cpu-offloading/benchmarks/v2v_prompt.json create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/config.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py create mode 100644 packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py create mode 100644 packages/diffusers-cpu-offloading/pyproject.toml diff --git a/README.md b/README.md index 6d3898dfa..f21583771 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ **Cosmos-Framework** is an end-to-end framework for training and serving world models, including the **Cosmos3** model family. Everything lives in a single top-level [`cosmos_framework/`](./cosmos_framework) Python package: - **Training** — distributed FSDP / TP / CP / PP trainer, native DCP checkpoints with HuggingFace `safetensors` import/export, JSONL / WebDataset / LeRobot dataset adapters. Entry point: `cosmos_framework.scripts.train`. See [`docs/training.md`](./docs/training.md). -- **Inference** — Diffusers / Transformers / vLLM backends with offline batch generation and online serving (Ray + Gradio). Entry point: `cosmos_framework.scripts.inference`. Ecosystem-facing shim libraries (lightweight standalone wrappers for downstream projects) live under [`packages/`](./packages). +- **Inference** — Diffusers / Transformers / vLLM backends with offline batch generation and online serving (Ray + Gradio). Entry point: `cosmos_framework.scripts.inference`. Ecosystem-facing shim libraries (lightweight standalone wrappers for downstream projects) live under [`packages/`](./packages). For Cosmos3-Nano generation on GPUs with limited VRAM, see the standalone [Diffusers CPU-GPU offloading adapter](./packages/diffusers-cpu-offloading/README.md). - **Policy serving** — a server-client pipeline for deploying Cosmos3 policies. See [`docs/action_policy_droid_server.md`](./docs/action_policy_droid_server.md). - **Post-training and RL** — supervised, preference-based, reinforcement-learning, distillation, diffusion, and embodied-policy optimization under `cosmos_framework.posttrain`. See [`docs/posttrain/`](./docs/posttrain/index.md). diff --git a/packages/diffusers-cpu-offloading/.gitignore b/packages/diffusers-cpu-offloading/.gitignore new file mode 100644 index 000000000..2e5ae9b6f --- /dev/null +++ b/packages/diffusers-cpu-offloading/.gitignore @@ -0,0 +1,11 @@ +.venv/ +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +artifacts/ +*.mp4 +*.jsonl +build/ +dist/ diff --git a/packages/diffusers-cpu-offloading/README.md b/packages/diffusers-cpu-offloading/README.md new file mode 100644 index 000000000..3b94cad16 --- /dev/null +++ b/packages/diffusers-cpu-offloading/README.md @@ -0,0 +1,193 @@ +# Cosmos3-Nano CPU-GPU weight streaming + +Run the Cosmos3-Nano Diffusers modular pipeline with model weights stored in +CPU RAM and streamed to the GPU as needed. Computation stays on the GPU. +The adapter supports BF16 and the official ModelOpt FP8 checkpoint without +modifying installed Diffusers source files. + +## Install into your existing environment + +Activate a compatible CUDA-enabled Python 3.12 environment. This standalone +package does not require installing the framework's training dependencies. From the +Cosmos-Framework repository root, change into this package and install the +adapter and benchmark dependencies: + +```bash +cd packages/diffusers-cpu-offloading +uv pip install '.[benchmarks]' +``` + +For FP8, install the optional dependency in the same environment: + +```bash +uv pip install '.[benchmarks,fp8]' +``` + +Run all examples and benchmark commands from this package directory. +The installer reads `pyproject.toml`; no `uv.lock` or `uv sync` is required. +You can use `python -m pip install` instead of `uv pip install`. +If you only need generation, install `.` for BF16 or `'.[fp8]'` for FP8. +Use `--no-deps` only when your environment already has all required dependencies. + +The adapter requires Python 3.12 or newer. Measurements used Linux, Python 3.12, +and PyTorch 2.13/CUDA 13. This package tracks Diffusers `main`; benchmark reports +record the installed commit, and upstream changes can affect compatibility and timings. +A BF16-capable NVIDIA GPU and substantial system RAM are required. The +consumer-GPU benchmark hosts had 128 GB RAM; that is a tested configuration, +not a measured minimum requirement for every workload. + +Accept the model and safety-model access terms on Hugging Face before running. +Use your normal Hugging Face authentication. Checkpoints are downloaded on +first use and reused from the cache afterward. Set `HF_HOME` to a sufficiently +large data volume if necessary. A manual checkpoint download is not required. + +## Generate video with BF16 + +```python +import torch +from diffusers.utils import export_to_video +from cosmos3_streaming import StreamingConfig, build_pipeline + +pipe = build_pipeline( + "nvidia/Cosmos3-Nano", + config=StreamingConfig(device="cuda:0"), + dtype=torch.bfloat16, +) + +frames = pipe( + prompt='{"scene":"A robot arm working in a bright kitchen"}', + negative_prompt="", + num_frames=189, + height=480, + width=848, + num_inference_steps=35, + guidance_scale=6.0, + fps=24.0, + generator=torch.Generator(device="cuda:0").manual_seed(0), + output="videos", +) +export_to_video(frames, "video.mp4", fps=24) +``` + +Use the official model-card prompts and conditioning assets for comparison with +the model's examples. The short prompt above illustrates the API. + +## Use the official FP8 checkpoint + +```python +pipe = build_pipeline( + "nvidia/Cosmos3-Nano", + revision="fp8", + checkpoint_precision="fp8", + config=StreamingConfig(device="cuda:0"), + dtype=torch.bfloat16, +) +``` + +FP8 uses the `fp8` branch of `nvidia/Cosmos3-Nano`. If you omit the model ID, +`checkpoint_precision="fp8"` selects that repository automatically. For this +repository, an omitted FP8 revision defaults to `fp8`; an explicit commit hash +or branch is preserved. Other repository IDs retain their own default branch. +BF16 continues to use `nvidia/Cosmos3-Nano` on its default branch. + +Diffusers reads the precision policy from the checkpoint and selects the +precision for each module and denoising step. Attention uses BF16 PyTorch SDPA. +The adapter refreshes cached understanding K/V when native precision changes. + +Native FP8 generation completed all 21 modality/resolution cases on RTX 5070, +with safety checking enabled except for explicitly labeled Transfer runs. +Use the benchmark runners below to validate performance and inspect outputs +on your hardware. + +## Other generation modalities + +Use the same `pipe` instance. Diffusers selects the generation workflow from +the call arguments. + +| Modality | Additional call inputs | +| --- | --- | +| Image-to-video | `image=` with a conditioning image | +| Video-to-video | `video=` and the appropriate frame-conditioning arguments | +| T2V or I2V with sound | `enable_sound=True` | +| Action | `action=CosmosActionCondition(...)` | +| Transfer | `control_videos=` with the matching control prompt and settings | + +For example, add a conditioning image to the video call: + +```python +from diffusers.utils import load_image + +image = load_image("conditioning.jpg") +frames = pipe( + prompt='{"scene":"The scene continues with natural motion"}', + negative_prompt="", + image=image, + num_frames=189, + height=480, + width=848, + num_inference_steps=35, + guidance_scale=6.0, + output="videos", +) +``` + +See the +[Cosmos3-Nano model card](https://huggingface.co/nvidia/Cosmos3-Nano) for official +inputs and generation settings. + +## Behavior and limitations + +- The default `memory_profile="auto"` chooses prefetch, VAE tiling, and the + constrained-memory Transfer strategy. Users do not need tile-size settings. +- Weight storage is in CPU RAM. Checkpoint files and the Hugging Face cache on + disk are not an implementation of per-block disk streaming. +- Safety checking is enabled by default. Rejections are not silently bypassed. + Any explicit use of `enable_safety_checker=False` remains the caller's + responsibility and must follow the model's applicable terms. +- Calls on a single pipeline instance must be sequential, not concurrent. + The instance can be reused for different supported workloads. +- The smaller GPU-memory footprint trades off against PCIe transfer time. + These helpers are for inference, not training. +- Cosmos3-Nano is the supported model. Do not assume Cosmos3-Super or every + action/control configuration has been validated. +- Runtime validation, failure cleanup, and optional diagnostic tracing remain + in the implementation. Experimental NVFP4 kernels and external kernel build + dependencies are not included. + +## Reproduce benchmarks + +Use the separate [BF16](benchmarks/benchmark_bf16.py) and +[FP8](benchmarks/benchmark_fp8.py) runners. See the +[benchmark guide](benchmarks/README.md) for setup, full modality/resolution +matrices, safety preflight, Transfer exceptions, and timing definitions. + +```bash +python benchmarks/benchmark_bf16.py --output-dir artifacts/benchmark-bf16 +python benchmarks/benchmark_fp8.py --output-dir artifacts/benchmark-fp8 +``` + +These run complete examples at each selected resolution. FP8 uses the +checkpoint-defined Diffusers policy and `nvidia/Cosmos3-Nano` with `revision="fp8"`. Reports +include generation settings, the precision policy, safety settings, and timings. + +## Source layout + +```text +cosmos3_streaming/ + __init__.py Public pipeline builders and configuration exports + config.py Runtime configuration + modular.py Diffusers modular-pipeline integration + runtime.py Conditioning cache and streamed execution + block_views.py Views of the original transformer blocks + offloading.py Group-offload lifecycle handling + fp8.py Official ModelOpt checkpoint restoration + attention/ BF16 SDPA implementation and backend interface + memory_trace.py Optional diagnostic tracing + metrics.py Lightweight request counters +benchmarks/ BF16/FP8 reproduction scripts, example inputs, and guide +``` + +This package includes the runtime adapter, BF16/FP8 benchmark runners, and small +benchmark input files. It does not include development tests, lab tooling, model +weights, or generated media. `pyproject.toml` supports installation directly from +this directory; no separately published package is required. diff --git a/packages/diffusers-cpu-offloading/benchmarks/README.md b/packages/diffusers-cpu-offloading/benchmarks/README.md new file mode 100644 index 000000000..1122b1579 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/README.md @@ -0,0 +1,210 @@ +# Reproduce Cosmos 3 Nano BF16 and FP8 benchmarks + +These runners generate complete examples using BF16 or the FP8 checkpoint's +native Diffusers policy, with the default Cosmos safety checker enabled. +Each run records generation settings, the precision policy, timings, and outputs. + +## Requirements + +- Linux, Python 3.12 and an NVIDIA GPU supported by your CUDA-enabled + PyTorch build, with a compatible NVIDIA driver. The target RTX 5070 has + 12 GB VRAM. We use a node configured with 128 GB system RAM; this is the test + configuration, not a measured minimum RAM requirement. +- Allow at least 80 GB per precision of disk space for dependencies, checkpoint caches, + compiler caches, and generated artifacts. A complete matrix can take hours. +- Access to `nvidia/Cosmos3-Nano` (including its `fp8` branch) and the safety-model repositories on Hugging + Face. Accept any applicable model terms. For gated downloads, supply an + `HF_TOKEN` through your usual secure environment setup. Do not put tokens + into scripts, logs, or a public issue. +- Use an otherwise idle GPU. CPU speed, PCIe bandwidth, GPU power limits, + driver version, and compiler-cache state can affect results. + +## Run + +From the Cosmos-Framework repository root, change into this package: + +```bash +cd packages/diffusers-cpu-offloading +``` + +Use your existing CUDA-enabled environment. Install the helper and benchmark +dependencies with `python -m pip install '.[benchmarks]'` for BF16, or +`python -m pip install '.[benchmarks,fp8]'` for FP8. An equivalent +`uv pip install` command works in an activated environment. No new environment +or `uv sync` is required. The installer provides the required dependencies, +and the runner checks the environment before generation. + +The package installs Diffusers from `main`. At the start of a benchmark, the +runner records the installed Git commit and package version in `recipe.json`. +Workers must use that revision, and resuming after a Diffusers upgrade requires +a new output directory. This keeps a matrix consistent without pinning every +user to one development commit. Compare reports from the same software revision +when evaluating performance differences. + +```bash +python benchmarks/benchmark_bf16.py --output-dir artifacts/benchmark-bf16 +python benchmarks/benchmark_fp8.py --output-dir artifacts/benchmark-fp8 +``` + +No manual checkpoint download or tiling settings are needed. BF16 uses +`nvidia/Cosmos3-Nano`; FP8 uses the same repository with `revision="fp8"`. +The runner resolves the branch to an immutable commit for each new run and +retains that commit on resume. The runner also +downloads the example prompts and conditioning inputs it needs. + +Subsequent runs reuse the dedicated cache under the output directory. +Use `--cache-dir /path/to/benchmark-cache` to share that cache between runs. +Use a **dedicated benchmark cache**, not your general Hugging Face cache: +the runner manages safety-model downloads there and runs generation offline. +Do not run concurrent preparations against the same cache. + +To run just one complete example first: + +```bash +python benchmarks/benchmark_bf16.py \ + --output-dir artifacts/bf16-t2v-480p \ + --modalities t2v --resolutions 480p +``` + +For FP8, use `benchmark_fp8.py` and a separate output directory with the same +modality and resolution arguments. Both commands use the full 35 denoising +steps and 189 frames. Safety checking cannot be disabled for this T2V example; +the explicit exception applies only to `--modalities transfer`. + +You can prepare downloads before occupying a GPU by adding `--prepare-only`. +Then repeat the same command with `--resume` instead of `--prepare-only`. +To resume an interrupted matrix, add `--resume` to the original command. Only +successful cases with verified outputs are skipped. When a case is retried, +its prior artifacts are preserved under `attempts/`. Use a new output directory +if you change the code or benchmark settings. + +## What the matrix runs + +| Modality flag | Example | Steps | Frames | FPS | Seed | Guidance | +|---|---|---:|---:|---:|---:|---:| +| `t2v` | Text to video | 35 | 189 | 24 | 123 | 6 | +| `i2v` | Image to video | 35 | 189 | 24 | 1111 | 6 | +| `t2vs` | Text to video with sound | 35 | 189 | 24 | 0 | 6 | +| `i2vs` | Image to video with sound | 35 | 189 | 24 | 0 | 6 | +| `action` | Action policy with video output | 30 | 17 | 5 | 0 | 1 | +| `v2v` | Video-conditioned generation | 35 | 189 | 24 | 0 | 6 | +| `transfer` | Edge-conditioned transfer | 35 per chunk | 121 | 30 | 0 | 3 | + +Each modality runs at `240p` (432x240), `480p` (848x480), and `720p` +(1280x720), in that order. Action uses its native resolution tiers 256, 480, +and 720 instead; the report records the actual output dimensions. Transfer +uses control guidance 1.5. All cases use UniPC, flow shift 10, no Karras sigmas, +and eight CPU compute threads. The adapter's automatic memory profile chooses +streaming and decode settings for the GPU; no tuning knobs are required. + +For the tested Action input, the 240p tier produces 320x240 and both the 480p +and 720p tiers produce 640x480. The 720p-tier Action result is not a 1280x720 +measurement. + +Cases run sequentially, each in a fresh process. This prevents one modality's +resident state from carrying into the next. The compiler cache persists across +cases. Within Transfer, the low-memory adapter stages guidance branches +according to its automatic memory policy. + +BF16 runs load the Nano checkpoint with `torch.bfloat16` and use BF16 PyTorch +SDPA attention. They do not use FP8 quantization or a mixed-precision schedule. + +For FP8, Diffusers resolves the checkpoint's precision policy, including the first/last +step counts and the precision assigned to each module group. The runner does +not hardcode those choices. A checkpoint without a mixed-precision schedule +uses its base W8A8 forwards. Cached understanding K/V is rebuilt when precision +changes, and the schedule restarts for each Transfer chunk. + +The FP8 audit uses read-only forward hooks to observe Diffusers' selected +precision during linear calls. It checks the selections against the resolved +checkpoint policy and requires observations for every denoising step. It does +not replace module forwards or execution functions, and it is not a hardware +kernel profiler. Reports contain the resolved policy and per-step observations. +Attention uses BF16 PyTorch SDPA. + +The default checker in `cosmos-guardrail==0.3.1` uses +Blocklist and Qwen3Guard for prompts, and RetinaFace face processing for video +output. Its separate video-content classifier is not enabled by the package +default. The runner preserves that default rather than claiming additional +checks. A rejected prompt or output is recorded as `safety_rejected`; it does +not silently disable checking, change the prompt, or invent a latency. + +## Read and compare results + +The output directory contains: + +```text +recipe.json # run configuration +prepared.json # download preparation details +summary.csv # table for analysis +summary.json +summary.md # readable consolidated table +completed.json # written only after the matrix finishes +t2v/240p/ + run.log + result.json # hardware, settings, timings, precision/safety audit + video.mp4 + memory_trace.jsonl +... # other modality/resolution cases +``` + +Sound cases also produce `audio.wav`; Action produces `action.json`. + +**Use `pipeline_seconds` to compare generation latency.** It measures the +complete pipeline call through CUDA completion: prompt checking, conditioning, +understanding, denoising, decoding, and output checking. Model loading, input +downloads, and MP4/WAV encoding are outside this timer. `load_seconds` and +`export_seconds` are recorded separately. `worker_seconds` covers the measured +worker body, including those stages and auditing, but excludes Python imports +and model/asset downloads. + +These are instrumented, single-run measurements, not warmed medians. Compilation +that occurs during the pipeline call is included. For variability estimates, +repeat into a new output directory with the same `--cache-dir`, then compare +the same cases and note that compiler caches are warm. Compare actual +dimensions, frames, seeds, precision and safety settings, not resolution labels +alone. A timing match does not establish pixel-identical outputs across GPUs; +inspect the videos as part of quality validation. + +The runner continues after failed cases and exits nonzero if any case did not +pass. `completed.json` means all requested cases were attempted, not that every +case passed. A `passed` result verifies execution, policy, safety calls, frame +count, and output geometry. It does not replace a human quality review. + +Video inputs come from NVIDIA's public examples. V2V prompt JSONs are bundled +with this runner. Transfer uses the Cosmos cookbook prompt. If the default guardrail +rejects it, the runner records the rejection. The explicit exception below +allows Transfer generation without checks while retaining that preflight result. + +## Prompt preflight and the Transfer exception + +Check prompts without downloading generation weights: + +```bash +python benchmarks/benchmark_bf16.py --output-dir artifacts/prompt-preflight \ + --preflight-only +``` + +Read `preflight.md` and `preflight.json` for the rejecting component and reason. +A rejection is distinct from a checker error. Use separate output directories +for preflight-only checks and generation. Both entry points +support preflight without downloading generation weights. + +Safety checks remain enabled for normal runs. To benchmark Transfer with +guardrails off, use a separate output directory: + +```bash +python benchmarks/benchmark_bf16.py --output-dir artifacts/bf16-transfer-no-guardrails \ + --modalities transfer --transfer-without-guardrails +python benchmarks/benchmark_fp8.py --output-dir artifacts/fp8-transfer-no-guardrails \ + --modalities transfer --transfer-without-guardrails +``` + +This option is rejected for any other modality selection. The prompt preflight +still runs and retains rejection evidence; checker errors are not bypassed. +Reports explicitly label checks as off. FP8 precision auditing is still required. +Use this exception only in accordance with the model's applicable terms. + +No benchmark is launched during installation. The scripts do not provision +cluster nodes or include internal lab paths. Run the commands on the GPU whose +performance you intend to measure. diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py new file mode 100644 index 000000000..e65baf08b --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py @@ -0,0 +1,6 @@ +"""Run full Cosmos3-Nano BF16 benchmarks with safety checking enabled by default.""" + +from runner import main + +if __name__ == "__main__": + raise SystemExit(main(precision="bf16")) diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py new file mode 100644 index 000000000..223f16f56 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py @@ -0,0 +1,6 @@ +"""Full Cosmos3-Nano FP8 benchmark using the checkpoint's native Diffusers policy.""" + +from runner import main + +if __name__ == "__main__": + raise SystemExit(main(precision="fp8")) diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py new file mode 100644 index 000000000..55915f496 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py @@ -0,0 +1,209 @@ +"""Shared, mandatory safety auditing for BF16 and FP8 benchmarks.""" + +import hashlib +import json +import time +from pathlib import Path + + +def benchmark_prompt(modality, model, inputs): + if modality == "action": + return "Put the pot to the left of the purple item." + if modality == "transfer": + path = Path(inputs) / "transfer/edge/prompt.json" + elif modality == "v2v": + path = Path(inputs) / "v2v/input_prompt.json" + else: + name = { + "t2v": "example_t2v_prompt.json", + "t2vs": "example_t2vs_prompt.json", + "i2v": "example_i2v_prompt.json", + "i2vs": "example_i2v_prompt.json", + }[modality] + path = Path(model) / "assets" / name + return json.dumps(json.loads(path.read_text())) + + +class GuardrailDisabledTransferAudit: + """Explicit Transfer exception; never report a safety check as passed.""" + + def __init__(self, pipe, precision="bf16"): + if precision not in ("bf16", "fp8"): + raise ValueError("Unsupported Transfer precision") + self.precision = precision + self.pipe = pipe + self.safety = [] + self._check_disabled() + + def _check_disabled(self): + if self.pipe.requires_safety_checker: + raise AssertionError("Transfer exception requires guardrails to be explicitly disabled") + + def validate(self, steps): + self._check_disabled() + count = self.pipe._streaming_runtime.stats.denoising_calls + if not count or count % steps: + raise AssertionError(f"Incomplete denoising: {count} calls for {steps} steps") + + def report(self): + return {"safety_checks": [], "guardrails": f"disabled_by_request_for_{self.precision}_transfer"} + + def close(self): + pass + + +class SafetyAudit: + """Observe the default checker, retaining its decisions and failing closed on errors.""" + + def __init__(self, pipe): + if not pipe.requires_safety_checker or pipe.safety_checker is None: + raise AssertionError("Safety checking is mandatory for benchmark results") + self.pipe = pipe + self.safety = [] + self.components = [] + self.originals = [] + checker = pipe.safety_checker + for model in checker.text_guardrail.safety_models: + original = model.is_safe + name = type(model).__name__ + self.originals.append((model, "is_safe", original)) + + def checked_component(prompt, _original=original, _name=name): + started = time.perf_counter() + try: + passed, reason = _original(prompt) + # Qwen3Guard 0.3.1 returns True on internal errors. Such an + # error is not a successful safety evaluation for a benchmark. + if "unexpected error" in reason.lower(): + raise RuntimeError(reason) + except Exception as exc: + self.components.append( + { + "component": _name, + "status": "error", + "reason": str(exc), + "seconds": time.perf_counter() - started, + } + ) + raise + self.components.append( + { + "component": _name, + "status": "passed" if passed else "rejected", + "reason": reason, + "seconds": time.perf_counter() - started, + } + ) + return passed, reason + + model.is_safe = checked_component + for kind, name in (("text", "check_text_safety"), ("video", "check_video_safety")): + original = getattr(checker, name) + self.originals.append((checker, name, original)) + + def checked(*args, _original=original, _kind=kind, **kwargs): + started = time.perf_counter() + try: + value = _original(*args, **kwargs) + except Exception as exc: + self.safety.append({"kind": _kind, "passed": False, "status": "error", "reason": str(exc)}) + raise + event = { + "kind": _kind, + "passed": bool(value) if _kind == "text" else value is not None, + "seconds": time.perf_counter() - started, + } + self.safety.append(event) + print(f"Safety: {event}", flush=True) + return value + + setattr(checker, name, checked) + + def report(self): + return {"safety_checks": self.safety, "prompt_components": self.components} + + def validate(self, steps=None): + if steps is not None: + count = self.pipe._streaming_runtime.stats.denoising_calls + if not count or count % steps: + raise AssertionError(f"Incomplete denoising: {count} calls for {steps} steps") + for kind in ("text", "video"): + events = [event for event in self.safety if event["kind"] == kind] + if not events or not all(event["passed"] for event in events): + raise AssertionError(f"Missing or rejected {kind} safety check") + + def close(self): + for obj, name, original in reversed(self.originals): + setattr(obj, name, original) + + +def run_preflight(model, inputs, modalities, resolutions, device, precision): + from types import SimpleNamespace + + import torch + from cosmos_guardrail.cosmos_guardrail import CosmosSafetyChecker + + # Load only the default guardrail models, never the Cosmos transformer, + # vision encoder or VAE. Use the same placement as the pipeline's prompt check. + checker = CosmosSafetyChecker() + checker.to(device) # cosmos-guardrail 0.3.1 .to() mutates in place and returns None + audit = SafetyAudit(SimpleNamespace(requires_safety_checker=True, safety_checker=checker)) + rows = [] + try: + for modality in modalities: + prompt = benchmark_prompt(modality, model, inputs) + audit.safety.clear() + audit.components.clear() + started = time.perf_counter() + error = None + try: + with torch.inference_mode(): + passed = checker.check_text_safety(prompt) + status = "passed" if passed else "safety_rejected" + except Exception as exc: # noqa: BLE001 - record errors separately, never as passes + status, error = "preflight_error", str(exc) + rows.append( + { + "modality": modality, + "resolutions": resolutions, + "precision": precision, + "status": status, + "stage": "before_generation", + "generation_started": False, + "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(), + "seconds": time.perf_counter() - started, + "error": error, + "safety_checks": list(audit.safety), + "components": list(audit.components), + } + ) + print(f"Preflight {modality}: {status}", flush=True) + finally: + audit.close() + checker.to("cpu") + return rows + + +def write_preflight_markdown(path, rows): + lines = [ + "# Prompt safety preflight", + "", + "These checks run before Cosmos model loading and generation. Output checks still run after generation.", + "", + "| Modality | Resolutions | Status | Rejecting/error component | Detail |", + "|---|---|---|---|---|", + ] + for row in rows: + failures = [event for event in row["components"] if event["status"] != "passed"] + names = ", ".join(event["component"] for event in failures) or "None" + reasons = [] + for event in failures: + reason = event["reason"] + # Keep large structured prompts out of the overview; full reasons + # remain in JSON and logs. Retain the exact matching keyword. + if "Match Word:" in reason: + reason = reason.split("Prompt:", 1)[0] + reason[reason.rfind("Match Word:") :] + reasons.append(reason.replace("|", "/").replace("\n", " ")) + detail = "; ".join(reasons) or row.get("error") or "Prompt checks passed" + lines.append(f"| {row['modality']} | {', '.join(row['resolutions'])} | {row['status']} | {names} | {detail} |") + Path(path).write_text("\n".join(lines) + "\n") diff --git a/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py b/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py new file mode 100644 index 000000000..32fac4576 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py @@ -0,0 +1,96 @@ +"""Observe Diffusers' checkpoint-defined FP8 policy during full generation.""" + +from collections import Counter +from dataclasses import asdict + +AUDIT_TYPE = "diffusers_fp8_policy" +AUDIT_SCHEMA_VERSION = 1 + + +class FP8PolicyAudit: + """Validate native precision selection without changing forwards or kernels.""" + + def __init__(self, pipe, safety_checker=True): + from benchmark_safety import GuardrailDisabledTransferAudit, SafetyAudit + from diffusers.pipelines.cosmos import mixed_precision as mixed + + self.mixed = mixed + self.pipe = pipe + self.policy = mixed.Cosmos3MixedPrecisionConfig.resolve(pipe.transformer) + self.groups = { + id(module): group + for name, module in pipe.transformer.named_modules() + if (group := mixed._classify_cosmos3_linear(name)) is not None and mixed._is_modelopt_fp8_linear(module) + } + if set(self.groups.values()) != {"generation", "reasoner"}: + raise AssertionError("Expected ModelOpt FP8 linears in both Diffusers precision groups") + self.safety_audit = ( + SafetyAudit(pipe) if safety_checker else GuardrailDisabledTransferAudit(pipe, precision="fp8") + ) + self.safety = self.safety_audit.safety + self.policy_validation_passed = False + self.selections = Counter() + self.handles = [] + try: + for module in pipe.transformer.modules(): + if id(module) in self.groups: + self.handles.append(module.register_forward_pre_hook(self._observe)) + except Exception: + self.close() + raise + + def _observe(self, module, _inputs): + step = self.pipe._streaming_runtime.stats.denoising_calls - 1 + group = self.groups[id(module)] + runtime = getattr(self.pipe.transformer, self.mixed._RUNTIME_ATTRIBUTE, None) + active = runtime is not None and runtime.active + if active and runtime.config != self.policy: + raise AssertionError("Diffusers runtime policy differs from the checkpoint policy") + if self.policy.enabled and not active: + raise AssertionError("Diffusers mixed-precision runtime is inactive during generation") + mode = "W8A16" if active and runtime.use_high_precision(group) else "W8A8" + self.selections[step, group, mode] += 1 + + def report(self): + return { + "audit_type": AUDIT_TYPE, + "schema_version": AUDIT_SCHEMA_VERSION, + "policy": asdict(self.policy), + "precision_dispatch": "native_diffusers", + "policy_validation_passed": self.policy_validation_passed, + **self.safety_audit.report(), + "observed_linears": [ + {"step": step, "group": group, "selected_precision": mode, "calls": count} + for (step, group, mode), count in sorted(self.selections.items()) + ], + } + + def validate(self, steps): + self.policy_validation_passed = False + total_steps = self.pipe._streaming_runtime.stats.denoising_calls + if steps <= 0 or not total_steps or total_steps % steps: + raise AssertionError(f"Incomplete denoising: {total_steps} calls for {steps} steps") + for step in range(total_steps): + if not sum(self.selections[step, "generation", mode] for mode in ("W8A8", "W8A16")): + raise AssertionError(f"Missing FP8 generation observations at step {step}") + if not any(group == "reasoner" for _, group, _ in self.selections): + raise AssertionError("Missing FP8 reasoner observations") + for (step, group, mode), count in self.selections.items(): + if not 0 <= step < total_steps: + raise AssertionError(f"FP8 observation outside the denoising loop: {step}") + high = self.policy.enabled and ( + self.policy.use_high_precision(step % steps, steps) + if group == "generation" + else self.policy.reasoner_policy == "high_precision" + ) + expected = "W8A16" if high else "W8A8" + if count and mode != expected: + raise AssertionError(f"Diffusers selected {mode}, expected {expected} for {group} at step {step}") + self.safety_audit.validate(steps) + self.policy_validation_passed = True + + def close(self): + for handle in self.handles: + handle.remove() + self.handles.clear() + self.safety_audit.close() diff --git a/packages/diffusers-cpu-offloading/benchmarks/manifest.json b/packages/diffusers-cpu-offloading/benchmarks/manifest.json new file mode 100644 index 000000000..6822b477e --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/manifest.json @@ -0,0 +1,94 @@ +{ + "schema_version": 2, + "bf16_model": { + "repo_id": "nvidia/Cosmos3-Nano", + "revision": "fea6e03ac3d7884b4105ed8ee79fc480fca70965" + }, + "name": "cosmos3-nano-native-streaming", + "model": { + "repo_id": "nvidia/Cosmos3-Nano", + "revision": "fp8" + }, + "safety_models": [ + { + "repo_id": "nvidia/Cosmos-1.0-Guardrail", + "revision": "cf03c0395fac8c4de386c0bdab12cc4fc8d66362", + "allow_patterns": [ + "blocklist/**", + "face_blur_filter/**" + ] + }, + { + "repo_id": "Qwen/Qwen3Guard-Gen-0.6B", + "revision": "fada3b2f655b89601929198343c94cd2f64d93cc", + "allow_patterns": [ + "*.json", + "*.safetensors", + "*.txt", + "*.jinja" + ] + } + ], + "assets": [ + { + "path": "v2v/input_robot_pouring.mp4", + "modalities": [ + "v2v" + ], + "url": "https://media.githubusercontent.com/media/nvidia-cosmos/cosmos-dependencies/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/vision/robot_pouring.mp4", + "sha256": "75b38640974fd5ad8245fa610acee9b0efe1f29d3c8fea829799d7b8fa799be8" + }, + { + "path": "action/bridge_20260501_0.mp4", + "modalities": [ + "action" + ], + "url": "https://media.githubusercontent.com/media/nvidia-cosmos/cosmos-dependencies/2b17a2413bd86b2cf9b03823637108851e4ddf2d/inputs/action/bridge_20260501_0.mp4", + "sha256": "a86cfc81633b216891ca26dc58c72193a979c10ad72f123175fa8d61a67cdaec" + }, + { + "path": "transfer/edge/control_edge.mp4", + "modalities": [ + "transfer" + ], + "url": "https://raw.githubusercontent.com/NVIDIA/cosmos/b0e54e88c322695dab188e6ed160c4d6d071c39d/cookbooks/cosmos3/generator/transfer/assets/edge/control_edge.mp4", + "sha256": "9d0fb1d1d29d5a081299a46d960fd496bce9997daecf4b39c16d2f4efa395719" + }, + { + "path": "transfer/edge/prompt.json", + "modalities": [ + "transfer" + ], + "url": "https://raw.githubusercontent.com/NVIDIA/cosmos/b0e54e88c322695dab188e6ed160c4d6d071c39d/cookbooks/cosmos3/generator/transfer/assets/edge/prompt.json", + "sha256": "03380ef96dceface2f4e79a8c526902d5654f556f38eebd7069a99f482655322" + }, + { + "path": "transfer/negative_prompt.json", + "modalities": [ + "transfer" + ], + "url": "https://raw.githubusercontent.com/NVIDIA/cosmos/b0e54e88c322695dab188e6ed160c4d6d071c39d/cookbooks/cosmos3/generator/transfer/assets/negative_prompt.json", + "sha256": "5ef32d9dcedfd5099be30d4dc9d5e959ad40da77f4d5daad782b40b195dbbb54" + }, + { + "path": "v2v/input_prompt.json", + "modalities": [ + "v2v" + ], + "bundled": "v2v_prompt.json", + "sha256": "e4721b5f27e73637c007ab24c8c6b513297df1d52708c4aa865fd085598b9c19" + }, + { + "path": "v2v/input_negative_prompt.json", + "modalities": [ + "v2v" + ], + "bundled": "v2v_negative_prompt.json", + "sha256": "7a60bd0c1335641cddd8b3ae66de767d2282f08861324f0ebe850e613af6a3a1" + } + ], + "example_model": { + "repo_id": "nvidia/Cosmos3-Nano", + "revision": "fea6e03ac3d7884b4105ed8ee79fc480fca70965" + } +} diff --git a/packages/diffusers-cpu-offloading/benchmarks/modalities.py b/packages/diffusers-cpu-offloading/benchmarks/modalities.py new file mode 100644 index 000000000..e308eefdd --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/modalities.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import Any + +import torch +from diffusers import CosmosActionCondition +from diffusers.utils import load_image, load_video + +RESOLUTIONS = { + "240p": (240, 432), + "480p": (480, 848), + "720p": (720, 1280), +} +ACTION_TIERS = { + "240p": 256, + "480p": 480, + "720p": 720, +} + + +def _load_json(path: Path) -> Any: + return json.loads(path.read_text()) + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _summarize_trace(path: Path) -> dict[str, Any]: + samples = [json.loads(line) for line in path.read_text().splitlines() if line] + return { + "samples": len(samples), + "peak_allocated_bytes": max((sample.get("global_peak_allocated_bytes", 0) for sample in samples), default=0), + "peak_reserved_bytes": max((sample.get("global_peak_reserved_bytes", 0) for sample in samples), default=0), + "peak_device_used_bytes": max( + (sample.get("global_peak_device_used_bytes", 0) for sample in samples), default=0 + ), + } + + +def _write_sound(result: dict[str, Any], path: Path) -> dict[str, Any]: + from scipy.io import wavfile + + waveform = result["sound"].detach().float().cpu().numpy() + if waveform.ndim == 3 and waveform.shape[0] == 1: + waveform = waveform[0] + if waveform.ndim != 2: + raise ValueError(f"expected Cosmos sound [channels, samples], got {tuple(waveform.shape)}") + sampling_rate = int(result["sampling_rate"]) + wavfile.write(path, sampling_rate, waveform.T) + return { + "path": str(path.resolve()), + "sha256": _sha256(path), + "channels": int(waveform.shape[0]), + "samples": int(waveform.shape[1]), + "sampling_rate": sampling_rate, + } + + +def _common_video_args(device: str, height: int, width: int, seed: int) -> dict[str, Any]: + return { + "num_frames": 189, + "height": height, + "width": width, + "fps": 24.0, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "add_resolution_template": False, + "add_duration_template": False, + "generator": torch.Generator(device=device).manual_seed(seed), + } + + +def _prepare_call( + modality: str, + model: Path, + device: str, + height: int, + width: int, + resolution: str, + v2v_assets: Path, + transfer_assets: Path, + action_video: Path, +) -> tuple[dict[str, Any], dict[str, Any], int]: + assets = model / "assets" + negative_prompt = _load_json(assets / "negative_prompt.json") + inputs: dict[str, Any] = { + "negative_prompt_json": str((assets / "negative_prompt.json").resolve()), + "negative_prompt_sha256": _sha256(assets / "negative_prompt.json"), + } + + if modality in {"t2v", "i2v", "t2vs", "i2vs"}: + prompt_name = ( + "example_t2vs_prompt.json" + if modality == "t2vs" + else ("example_i2v_prompt.json" if modality in {"i2v", "i2vs"} else "example_t2v_prompt.json") + ) + prompt_path = assets / prompt_name + call = _common_video_args(device, height, width, seed=0 if modality.endswith("s") else 123) + if modality == "i2v": + call["generator"] = torch.Generator(device=device).manual_seed(1111) + call.update( + { + "prompt": json.dumps(_load_json(prompt_path)), + "negative_prompt": json.dumps(negative_prompt), + } + ) + inputs.update({"prompt_json": str(prompt_path.resolve()), "prompt_sha256": _sha256(prompt_path)}) + if modality in {"i2v", "i2vs"}: + image_path = assets / "example_i2v_input.jpg" + call["image"] = load_image(str(image_path)) + inputs.update({"image": str(image_path.resolve()), "image_sha256": _sha256(image_path)}) + if modality in {"t2vs", "i2vs"}: + call["enable_sound"] = True + call["output"] = ["videos", "sound", "sampling_rate"] + else: + call["output"] = "videos" + return call, inputs, 24 + + if modality == "v2v": + video_path = v2v_assets / "input_robot_pouring.mp4" + prompt_path = v2v_assets / "input_prompt.json" + negative_path = v2v_assets / "input_negative_prompt.json" + call = _common_video_args(device, height, width, seed=0) + call.update( + { + "prompt": json.dumps(_load_json(prompt_path)), + "negative_prompt": json.dumps(_load_json(negative_path)), + "video": load_video(str(video_path)), + "condition_frame_indexes_vision": [0, 1], + "condition_video_keep": "first", + "output": "videos", + } + ) + inputs = { + "video": str(video_path.resolve()), + "video_sha256": _sha256(video_path), + "prompt_json": str(prompt_path.resolve()), + "prompt_sha256": _sha256(prompt_path), + "negative_prompt_json": str(negative_path.resolve()), + "negative_prompt_sha256": _sha256(negative_path), + } + return call, inputs, 24 + + if modality == "transfer": + video_path = transfer_assets / "edge" / "control_edge.mp4" + prompt_path = transfer_assets / "edge" / "prompt.json" + negative_path = transfer_assets / "negative_prompt.json" + call = { + "prompt": json.dumps(_load_json(prompt_path)), + "negative_prompt": json.dumps(_load_json(negative_path)), + "control_videos": {"edge": load_video(str(video_path))}, + "num_frames": 121, + "height": height, + "width": width, + "fps": 30.0, + "num_inference_steps": 35, + "guidance_scale": 3.0, + "control_guidance": 1.5, + "generator": torch.Generator(device=device).manual_seed(0), + "output": "videos", + } + inputs = { + "control_video": str(video_path.resolve()), + "control_video_sha256": _sha256(video_path), + "prompt_json": str(prompt_path.resolve()), + "prompt_sha256": _sha256(prompt_path), + "negative_prompt_json": str(negative_path.resolve()), + "negative_prompt_sha256": _sha256(negative_path), + } + return call, inputs, 30 + + if modality == "action": + tier = ACTION_TIERS[resolution] + call = { + "prompt": "Put the pot to the left of the purple item.", + "action": CosmosActionCondition( + mode="policy", + chunk_size=16, + domain_name="bridge_orig_lerobot", + resolution_tier=tier, + video=load_video(str(action_video)), + view_point="ego_view", + ), + "fps": 5, + "num_inference_steps": 30, + "guidance_scale": 1.0, + "use_system_prompt": False, + "generator": torch.Generator(device=device).manual_seed(0), + "output": ["videos", "action"], + } + inputs = {"video": str(action_video.resolve()), "video_sha256": _sha256(action_video)} + return call, inputs, 5 + + raise AssertionError(f"unhandled modality: {modality}") diff --git a/packages/diffusers-cpu-offloading/benchmarks/runner.py b/packages/diffusers-cpu-offloading/benchmarks/runner.py new file mode 100644 index 000000000..6c6f0bbf0 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/runner.py @@ -0,0 +1,630 @@ +"""Run full BF16 or checkpoint-policy FP8 benchmarks with safety checking enabled. + +Default: seven modalities, three resolutions, full steps, one GPU, fresh process +per case. See benchmarks/README.md for timing definitions and requirements. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import importlib.metadata +import json +import os +import shutil +import subprocess +import sys +import time +import traceback +import urllib.request +from pathlib import Path + +from fp8_benchmark_audit import AUDIT_SCHEMA_VERSION, AUDIT_TYPE + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +MANIFEST = Path(__file__).with_name("manifest.json") +MODALITIES = ("t2v", "i2v", "t2vs", "i2vs", "action", "v2v", "transfer") +RESOLUTIONS = ("240p", "480p", "720p") + + +def diffusers_installation(): + """Identify the installed Git revision without importing the GPU runtime.""" + distribution = importlib.metadata.distribution("diffusers") + origin = json.loads(distribution.read_text("direct_url.json") or "{}") + commit = origin.get("vcs_info", {}).get("commit_id") + if not isinstance(commit, str) or len(commit) != 40 or any(c not in "0123456789abcdef" for c in commit): + raise ValueError( + "Benchmarks require a Git installation of Diffusers with an immutable commit ID. " + "Install this package's dependencies from pyproject.toml; see benchmarks/README.md." + ) + return {"diffusers_commit": commit, "diffusers_version": distribution.version} + + +def check_diffusers_installation(manifest): + """Prevent workers from mixing revisions within a prepared benchmark.""" + installed = diffusers_installation() + if any(manifest.get(key) != value for key, value in installed.items()): + raise ValueError("Diffusers changed after benchmark preparation. Use a new output directory.") + + +def digest(path): + with Path(path).open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def write_json(path, value): + path = Path(path) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2) + "\n") + temporary.replace(path) + + +def fingerprint(): + paths = [ + *ROOT.glob("cosmos3_streaming/**/*.py"), + *ROOT.glob("benchmarks/**/*.py"), + *ROOT.glob("benchmarks/**/*.json"), + ROOT / "pyproject.toml", + ] + files = {str(path.relative_to(ROOT)): digest(path) for path in sorted(paths) if path.is_file()} + return {"sha256": hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest(), "files": files} + + +def configure_cache(cache): + cache = cache.resolve() + paths = { + "HF_HOME": cache / "huggingface", + "HF_HUB_CACHE": cache / "huggingface/hub", + "HUGGINGFACE_HUB_CACHE": cache / "huggingface/hub", + "XDG_CACHE_HOME": cache, + "TORCH_HOME": cache / "torch", + "TRITON_CACHE_DIR": cache / "triton", + "TORCHINDUCTOR_CACHE_DIR": cache / "inductor", + "CUDA_CACHE_PATH": cache / "cuda", + "TMPDIR": cache / "tmp", + "XDG_CONFIG_HOME": cache / "config", + "XDG_DATA_HOME": cache / "data", + } + for name, path in paths.items(): + path.mkdir(parents=True, exist_ok=True) + os.environ[name] = str(path) + # Fixed CPU threading is part of the benchmark recipe. + for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"): + os.environ[name] = "8" + os.environ["PYTHONUNBUFFERED"] = "1" + + +def prepare(args, manifest): + from huggingface_hub import snapshot_download + + assets = args.output_dir / "inputs" + for item in manifest["assets"]: + if not set(args.modalities).intersection(item["modalities"]): + continue + target = assets / item["path"] + target.parent.mkdir(parents=True, exist_ok=True) + if not target.exists(): + temporary = target.with_suffix(target.suffix + ".download") + if "bundled" in item: + shutil.copyfile(MANIFEST.parent / item["bundled"], temporary) + else: + print(f"Downloading {item['path']}", flush=True) + with urllib.request.urlopen(item["url"], timeout=120) as source, temporary.open("wb") as dest: + shutil.copyfileobj(source, dest) + if digest(temporary) != item["sha256"]: + raise ValueError(f"Asset checksum mismatch: {item['path']}") + temporary.replace(target) + if digest(target) != item["sha256"]: + raise ValueError(f"Existing asset checksum mismatch: {target}") + + examples = Path(snapshot_download(**manifest["example_model"], allow_patterns=["assets/**"])) + model = None + if not args.preflight_only: + model = Path( + snapshot_download( + repo_id=manifest["model"]["repo_id"], + revision=manifest["model"]["revision"], + allow_patterns=[ + "*.json", + "*.txt", + "*.model", + "*.jinja", + "transformer/**", + "vae/**", + "vision_encoder/**", + "sound_tokenizer/**", + "text_tokenizer/**", + "scheduler/**", + ], + ) + ) + for safety in manifest["safety_models"]: + snapshot = Path(snapshot_download(**safety)) + # Guardrail 0.3.1 requests main internally. Pin it in this benchmark's + # dedicated cache; workers are offline, so it cannot drift during a run. + refs = snapshot.parent.parent / "refs" + refs.mkdir(exist_ok=True) + (refs / "main").write_text(safety["revision"]) + model_hashes = { + str(path.relative_to(model)): digest(path) + for path in (sorted(model.rglob("*")) if model is not None else []) + if path.is_file() and ".cache" not in path.relative_to(model).parts + } + prepared = { + "model_path": str(model) if model is not None else None, + "example_model_path": str(examples), + "example_files_sha256": { + str(path.relative_to(examples)): digest(path) + for path in sorted((examples / "assets").rglob("*")) + if path.is_file() + }, + "manifest": manifest, + "model_files_sha256": model_hashes, + "weights_requested": not args.preflight_only, + } + write_json(args.output_dir / "prepared.json", prepared) + return prepared + + +def resolve_model(manifest, previous=None): + """Pin mutable Hub revisions once; a resumed run retains its original commit.""" + from huggingface_hub import HfApi + + requested = manifest["requested_model"] + selected = dict(requested) + if previous and previous["manifest"].get("requested_model") == requested: + selected = dict(previous["manifest"]["model"]) + info = HfApi().model_info(selected["repo_id"], revision=selected["revision"]) + names = {item.rfilename for item in info.siblings} + if "transformer/config.json" not in names or not any( + name.startswith("transformer/") and name.endswith((".safetensors", ".bin", ".pt", ".pth")) for name in names + ): + raise ValueError( + f"{selected['repo_id']}@{selected['revision']} does not contain a transformer checkpoint. " + "Wait for the weights to be published or select a published --revision. " + "No fallback to a different checkpoint or precision is performed." + ) + selected["revision"] = info.sha + return selected + + +def guardrails_enabled(manifest, modality): + exception = manifest.get("transfer_without_guardrails", False) + if exception and manifest.get("checkpoint_precision") not in ("bf16", "fp8"): + raise ValueError("The guardrail exception requires BF16 or FP8 Transfer") + return not (exception and modality == "transfer") + + +def preflight_allows_generation(row, enabled): + # Explicit exceptions may bypass a rejection, never a checker/infrastructure error. + return row["status"] == "passed" or (not enabled and row["status"] == "safety_rejected") + + +def passed_and_intact(report, safety_checker=True): + checks = report.get("audit", {}).get("safety_checks", []) + if safety_checker: + safety_valid = all( + any(event.get("kind") == kind and event.get("passed") is True for event in checks) + for kind in ("text", "video") + ) + else: + precision = report.get("checkpoint_precision") + safety_valid = ( + report.get("modality") == "transfer" + and precision in ("bf16", "fp8") + and report.get("audit", {}).get("guardrails") == f"disabled_by_request_for_{precision}_transfer" + and (precision != "fp8" or report.get("audit", {}).get("policy_validation_passed") is True) + and not checks + ) + return ( + report.get("status") == "passed" + and ( + report.get("checkpoint_precision") != "fp8" + or ( + report.get("audit", {}).get("precision_dispatch") == "native_diffusers" + and report.get("audit", {}).get("policy_validation_passed") is True + and report.get("audit", {}).get("audit_type") == AUDIT_TYPE + and report.get("audit", {}).get("schema_version") == AUDIT_SCHEMA_VERSION + ) + ) + and report.get("safety_checker") is safety_checker + and safety_valid + and bool(report.get("artifacts")) + and all(Path(item["path"]).is_file() and digest(item["path"]) == item["sha256"] for item in report["artifacts"]) + ) + + +def summarize(root, cases, precision="fp8"): + rows = [] + for modality, resolution in cases: + report_path = root / modality / resolution / "result.json" + report = json.loads(report_path.read_text()) if report_path.exists() else {} + rows.append( + { + "modality": modality, + "resolution": resolution, + "precision": "BF16" if precision == "bf16" else "FP8 (checkpoint policy)", + "fp8_policy": report.get("audit", {}).get("policy") if precision == "fp8" else None, + "safety_checker": report.get("safety_checker"), + "status": report.get("status", "pending"), + "gpu": report.get("system", {}).get("gpu", ""), + "output_width": report.get("outputs", {}).get("width", ""), + "output_height": report.get("outputs", {}).get("height", ""), + "frames": report.get("outputs", {}).get("frames", ""), + **{ + key: report.get("timing", {}).get(key, "") + for key in ("load_seconds", "pipeline_seconds", "export_seconds", "worker_seconds") + }, + } + ) + write_json(root / "summary.json", rows) + with (root / "summary.csv").open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + lines = [ + "| Modality | Resolution | Output | Frames | Pipeline (s) | Guardrails | Status |", + "|---|---|---|---:|---:|---|---|", + ] + for row in rows: + seconds = row["pipeline_seconds"] + timing = f"{seconds:.3f}" if isinstance(seconds, (int, float)) else "N/A" + safety = {True: "On", False: "Off", None: "Not recorded"}[row["safety_checker"]] + lines.append( + f"| {row['modality']} | {row['resolution']} | " + f"{row['output_width']}x{row['output_height']} | {row['frames']} | {timing} | {safety} | {row['status']} |" + ) + (root / "summary.md").write_text("\n".join(lines) + "\n") + return rows + + +def worker(args): + # Only workers import CUDA libraries. Downloads and matrix orchestration + # never create a CUDA context and are excluded from pipeline latency. + from dataclasses import asdict + + import torch + from benchmark_safety import GuardrailDisabledTransferAudit, SafetyAudit, benchmark_prompt + from diffusers.schedulers import UniPCMultistepScheduler + from diffusers.utils import export_to_video + from fp8_benchmark_audit import FP8PolicyAudit + from modalities import RESOLUTIONS as SIZES + from modalities import _prepare_call, _summarize_trace, _write_sound + + from cosmos3_streaming import StreamingConfig, build_pipeline + + modality, resolution = args.worker + directory = args.output_dir / modality / resolution + directory.mkdir(parents=True, exist_ok=True) + prepared = json.loads((args.output_dir / "prepared.json").read_text()) + precision = prepared["manifest"].get("checkpoint_precision", "fp8") + safety_checker = guardrails_enabled(prepared["manifest"], modality) + trace = directory / "memory_trace.jsonl" + report = { + "status": "running", + "modality": modality, + "resolution": resolution, + "full_generation": True, + "checkpoint_precision": precision, + "safety_checker": safety_checker, + "timing": {}, + "artifacts": [], + } + audit = None + if not safety_checker: + report["safety_exception"] = f"Explicit {precision.upper()} Transfer run with guardrails off" + report["prompt_preflight"] = next( + row for row in json.loads((args.output_dir / "preflight.json").read_text()) if row["modality"] == modality + ) + started = time.perf_counter() + try: + torch.cuda.set_device(args.device) + properties = torch.cuda.get_device_properties(args.device) + report["system"] = { + "gpu": properties.name, + "vram_bytes": properties.total_memory, + "compute_capability": f"{properties.major}.{properties.minor}", + "host": os.uname().nodename, + "python": sys.version, + "cuda": torch.version.cuda, + "ram_bytes": os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES"), + "versions": { + name: importlib.metadata.version(name) + for name in ( + ("torch", "diffusers", "transformers", "cosmos-guardrail") + + (("nvidia-modelopt",) if precision == "fp8" else ()) + ) + }, + "nvidia_smi": subprocess.check_output( + [ + "nvidia-smi", + ( + "--query-gpu=name,uuid,memory.total,driver_version," + "power.limit,pcie.link.gen.max,pcie.link.width.max" + ), + "--format=csv", + ], + text=True, + ), + } + origin = json.loads(importlib.metadata.distribution("diffusers").read_text("direct_url.json") or "{}") + report["diffusers_origin"] = origin + check_diffusers_installation(prepared["manifest"]) + height, width = SIZES[resolution] + inputs = args.output_dir / "inputs" + call, report["inputs"], fps = _prepare_call( + modality, + Path(prepared["example_model_path"]), + args.device, + height, + width, + resolution, + inputs / "v2v", + inputs / "transfer", + inputs / "action/bridge_20260501_0.mp4", + ) + if call["prompt"] != benchmark_prompt(modality, prepared["example_model_path"], inputs): + raise AssertionError("Generation prompt differs from the preflight prompt") + report["settings"] = { + key: value + for key, value in call.items() + if isinstance(value, (str, int, float, bool)) or key in ("output", "condition_frame_indexes_vision") + } + report["settings"]["seed"] = call["generator"].initial_seed() + report["settings"]["flow_shift"] = 10.0 + report["settings"]["use_karras_sigmas"] = False + if modality == "action": + report["settings"]["action_resolution_tier"] = {"240p": 256, "480p": 480, "720p": 720}[resolution] + config = StreamingConfig(device=args.device, memory_trace_path=str(trace), memory_trace_blocks=True) + report["streaming_config"] = asdict(config) + load_start = time.perf_counter() + pipe = build_pipeline( + prepared["model_path"], + config=config, + dtype=torch.bfloat16, + checkpoint_precision=precision, + local_files_only=True, + enable_safety_checker=safety_checker, + ) + pipe.scheduler = UniPCMultistepScheduler.from_config( + pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False + ) + torch.cuda.synchronize(args.device) + report["timing"]["load_seconds"] = time.perf_counter() - load_start + if precision == "fp8": + audit = FP8PolicyAudit(pipe, safety_checker=safety_checker) + elif not safety_checker: + audit = GuardrailDisabledTransferAudit(pipe) + else: + audit = SafetyAudit(pipe) + write_json(directory / "result.json", report) + print(f"Generating full {modality} {resolution} on {properties.name}", flush=True) + generation_start = time.perf_counter() + result = pipe(**call) + torch.cuda.synchronize(args.device) + report["timing"]["pipeline_seconds"] = time.perf_counter() - generation_start + audit.validate(call["num_inference_steps"]) + videos = result["videos"] if isinstance(result, dict) else result + expected_frames = 17 if modality == "action" else call["num_frames"] + if len(videos) != expected_frames: + raise AssertionError(f"Expected {expected_frames} frames, received {len(videos)}") + output_width, output_height = videos[0].size + if modality != "action" and (output_height, output_width) != (height, width): + raise AssertionError("Output resolution does not match the requested resolution") + export_start = time.perf_counter() + video = directory / "video.mp4" + export_to_video(videos, str(video), fps=fps, macro_block_size=1) + outputs = {"frames": len(videos), "width": output_width, "height": output_height, "fps": fps} + artifacts = [video] + if modality in ("t2vs", "i2vs"): + sound = directory / "audio.wav" + outputs["sound"] = _write_sound(result, sound) + artifacts.append(sound) + if modality == "action": + action = directory / "action.json" + write_json(action, result["action"][0].float().cpu().tolist()) + artifacts.append(action) + report["timing"]["export_seconds"] = time.perf_counter() - export_start + report["outputs"] = outputs + report["artifacts"] = [{"path": str(path.resolve()), "sha256": digest(path)} for path in artifacts] + report["runtime_stats"] = asdict(pipe._streaming_runtime.stats) + report["status"] = "passed" + except Exception: # noqa: BLE001 - persist any failed case and continue the matrix + rejected = audit is not None and any( + not event["passed"] and event.get("status") != "error" for event in audit.safety + ) + report["status"] = "safety_rejected" if rejected else "failed" + report["error"] = traceback.format_exc() + # A failed/rejected case has no valid generation latency. + report["timing"].pop("pipeline_seconds", None) + print(report["error"], flush=True) + finally: + if audit is not None: + report["audit"] = audit.report() + audit.close() + if trace.exists(): + report["memory"] = _summarize_trace(trace) + report["timing"]["worker_seconds"] = time.perf_counter() - started + write_json(directory / "result.json", report) + return 0 if report["status"] == "passed" else 1 + + +def main(precision="fp8"): + parser = argparse.ArgumentParser(description=f"Full {precision.upper()} benchmark with safety checking by default") + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--cache-dir", type=Path, help="Dedicated benchmark cache (default: OUTPUT/cache)") + parser.add_argument("--modalities", nargs="+", choices=MODALITIES, default=list(MODALITIES)) + parser.add_argument("--resolutions", nargs="+", choices=RESOLUTIONS, default=list(RESOLUTIONS)) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--prepare-only", action="store_true") + parser.add_argument( + "--preflight-only", + action="store_true", + help="Check every selected prompt before downloading generation weights or generating media", + ) + parser.add_argument("--resume", action="store_true") + parser.add_argument( + "--transfer-without-guardrails", + action="store_true", + help="Explicit Transfer-only exception; retain prompt rejection evidence but generate with guardrails off", + ) + parser.add_argument("--revision", help="Override the checkpoint revision; resolved to an immutable Hub commit") + parser.add_argument("--worker", nargs=2, metavar=("MODALITY", "RESOLUTION"), help=argparse.SUPPRESS) + parser.add_argument("--preflight-worker", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + if args.transfer_without_guardrails and args.modalities != ["transfer"]: + parser.error("--transfer-without-guardrails requires --modalities transfer") + args.output_dir = args.output_dir.resolve() + args.output_dir.mkdir(parents=True, exist_ok=True) + configure_cache(args.cache_dir or args.output_dir / "cache") + if args.worker: + return worker(args) + if args.preflight_worker: + from benchmark_safety import run_preflight, write_preflight_markdown + + prepared = json.loads((args.output_dir / "prepared.json").read_text()) + try: + check_diffusers_installation(prepared["manifest"]) + rows = run_preflight( + prepared["example_model_path"], + args.output_dir / "inputs", + args.modalities, + args.resolutions, + args.device, + prepared["manifest"]["checkpoint_precision"], + ) + except Exception as exc: # noqa: BLE001 - initialization failures must not disappear from the report + rows = [ + { + "modality": modality, + "resolutions": args.resolutions, + "status": "preflight_error", + "stage": "preflight_initialization", + "generation_started": False, + "components": [], + "error": str(exc), + } + for modality in args.modalities + ] + traceback.print_exc() + write_json(args.output_dir / "preflight.json", rows) + write_preflight_markdown(args.output_dir / "preflight.md", rows) + return 0 # Rejection/error statuses are in the report and handled by the parent. + manifest = json.loads(MANIFEST.read_text()) + manifest.update(diffusers_installation()) + if precision == "bf16": + manifest["model"] = manifest["bf16_model"] + manifest["name"] = "cosmos3-nano-bf16-safety" + if args.revision: + manifest["model"]["revision"] = args.revision + manifest["checkpoint_precision"] = precision + if args.transfer_without_guardrails: + manifest["transfer_without_guardrails"] = True + recipe = args.output_dir / "recipe.json" + previous = json.loads(recipe.read_text()) if recipe.exists() else None + manifest["requested_model"] = dict(manifest["model"]) + if not args.preflight_only: + manifest["model"] = resolve_model(manifest, previous) + identity = { + "source": fingerprint(), + "preflight_only": args.preflight_only, + "device": args.device, + "manifest": manifest, + "modalities": args.modalities, + "resolutions": args.resolutions, + } + if recipe.exists(): + if previous != identity: + raise ValueError("Recipe/source changed. Use a new output directory, never mix benchmark results") + if not args.resume: + raise ValueError("Output already initialized. Use --resume or a new output directory") + write_json(recipe, identity) + prepare(args, manifest) + if args.prepare_only: + print(f"Prepared checkpoint and inputs in {args.output_dir}") + return 0 + cases = [(modality, resolution) for modality in args.modalities for resolution in args.resolutions] + summarize(args.output_dir, cases, precision) + env = dict(os.environ, HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1") + preflight_command = [ + sys.executable, + str(Path(__file__).resolve()), + "--output-dir", + str(args.output_dir), + "--cache-dir", + os.environ["XDG_CACHE_HOME"], + "--device", + args.device, + "--preflight-worker", + "--modalities", + *args.modalities, + "--resolutions", + *args.resolutions, + ] + print(f"Checking prompts before generation; log: {args.output_dir / 'preflight.log'}", flush=True) + with (args.output_dir / "preflight.log").open("w") as log: + subprocess.run(preflight_command, env=env, stdout=log, stderr=subprocess.STDOUT, check=True) + preflight = {row["modality"]: row for row in json.loads((args.output_dir / "preflight.json").read_text())} + if args.preflight_only: + print((args.output_dir / "preflight.md").read_text(), flush=True) + return 0 if all(row["status"] == "passed" for row in preflight.values()) else 1 + for modality, resolution in cases: + safety_checker = guardrails_enabled(manifest, modality) + directory = args.output_dir / modality / resolution + result_path = directory / "result.json" + if result_path.exists() and passed_and_intact(json.loads(result_path.read_text()), safety_checker): + print(f"Verified completed case: {modality}/{resolution}", flush=True) + continue + if directory.exists(): + archive = args.output_dir / "attempts" / f"{modality}-{resolution}-{time.time_ns()}" + archive.parent.mkdir(exist_ok=True) + directory.rename(archive) + directory.mkdir(parents=True) + if not preflight_allows_generation(preflight[modality], safety_checker): + write_json( + result_path, + { + "status": preflight[modality]["status"], + "modality": modality, + "resolution": resolution, + "checkpoint_precision": precision, + "safety_checker": True, + "generation_started": False, + "failure_stage": "prompt_preflight", + "preflight": preflight[modality], + }, + ) + summarize(args.output_dir, cases, precision) + continue + command = [ + sys.executable, + str(Path(__file__).resolve()), + "--output-dir", + str(args.output_dir), + "--cache-dir", + os.environ["XDG_CACHE_HOME"], + "--device", + args.device, + "--worker", + modality, + resolution, + ] + print(f"Starting {modality}/{resolution}; log: {directory / 'run.log'}", flush=True) + with (directory / "run.log").open("w") as log: + completed = subprocess.run(command, env=env, stdout=log, stderr=subprocess.STDOUT, check=False) + if not result_path.exists() or json.loads(result_path.read_text()).get("status") == "running": + write_json(result_path, {"status": "failed", "error": f"Worker exited {completed.returncode}; see run.log"}) + print(f"Finished {modality}/{resolution}: {json.loads(result_path.read_text())['status']}", flush=True) + summarize(args.output_dir, cases, precision) + rows = summarize(args.output_dir, cases, precision) + write_json( + args.output_dir / "completed.json", + {"cases": len(rows), "passed": sum(row["status"] == "passed" for row in rows), "finished_at_unix": time.time()}, + ) + return 0 if all(row["status"] == "passed" for row in rows) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/diffusers-cpu-offloading/benchmarks/v2v_negative_prompt.json b/packages/diffusers-cpu-offloading/benchmarks/v2v_negative_prompt.json new file mode 100644 index 000000000..2a0bd29b7 --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/v2v_negative_prompt.json @@ -0,0 +1,108 @@ +{ + "subjects": [ + { + "description": "Blurry, poorly defined subjects with inconsistent shapes and unrealistic proportions.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Extremely low-quality subjects with visible rendering artifacts, broken mesh geometry, and completely unrealistic proportions throughout.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Poorly generated subjects exhibiting all hallmarks of failed neural rendering \u2014 flickering edges, inconsistent depth, and uncanny spatial relationships.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + } + ], + "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity \u2014 patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.", + "lighting": { + "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.", + "direction": "Inconsistent light sources \u2014 shadows point in multiple contradictory directions, breaking physical plausibility.", + "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable \u2014 even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.", + "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s." + }, + "aesthetics": { + "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.", + "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.", + "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.", + "patterns": "Visible tiling artifacts in textures, moir\u00e9 patterns, and aliasing on edges." + }, + "cinematography": { + "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow.", + "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space.", + "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content \u2014 panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles \u2014 instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.", + "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.", + "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.", + "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression." + }, + "style_medium": "Low quality compressed digital video with visible encoding artifacts", + "artistic_style": "Amateur, unpolished with inconsistent visual style", + "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.", + "actions": [ + { + "time": "0:00-0:08", + "description": "Subjects attempt to move but their motion is jerky, temporally inconsistent, and physically implausible. Background elements flicker and shift between frames." + } + ], + "text_and_signage_elements": [], + "segments": [ + { + "segment_index": 0, + "time_range": "0:00-0:08", + "description": "A single continuous shot suffering from severe temporal inconsistencies \u2014 subjects that morph and deform between frames, backgrounds that shift and wobble, and rendering quality that fluctuates visibly over time. Motion blur is applied incorrectly, smearing in directions that don't match actual movement. Frame-to-frame coherence breaks down with individual pixels changing color randomly in flat areas. Texture detail level fluctuates between frames as if the rendering budget varied shot to shot. Color grading drifts over the duration with no creative motivation. Noise patterns change between frames in ways that draw attention rather than being invisible. Overall visual quality degrades progressively from start to finish.", + "key_changes": "No meaningful progression or narrative development. Visual quality degrades over time.", + "camera": "Unstable, poorly controlled camera work with visible mathematical interpolation artifacts." + } + ], + "transitions": [], + "temporal_caption": "The scene opens at 0.0 seconds with a poorly rendered establishing shot that immediately reveals low production quality. At 1.0 seconds, subjects begin to move but their motion is jerky and inconsistent, with limbs bending at unnatural angles and objects clipping through each other. From 2.0 to 4.0 seconds, the camera shakes violently while the scene exhibits visible compression artifacts, color banding in the sky, and flickering in the shadows. Between 4.0 and 6.0 seconds, temporal coherence breaks down as elements appear and disappear between frames, textures swim and morph unnaturally, and the lighting shifts abruptly without physical cause. In the final 2 seconds, the overall visual quality deteriorates further with increasing noise, blur, and a general loss of spatial coherence that makes the scene nearly unwatchable. Additionally, the frame rate appears inconsistent with visible judder and stuttering throughout. Color temperature shifts randomly between warm and cool tones with no motivation. The encode quality degrades in complex regions showing macro-blocking and mosquito noise around moving edges. Temporal noise patterns are spatially correlated, creating swimming artifacts on flat surfaces.", + "audio_description": "", + "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent \u2014 objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded." +} diff --git a/packages/diffusers-cpu-offloading/benchmarks/v2v_prompt.json b/packages/diffusers-cpu-offloading/benchmarks/v2v_prompt.json new file mode 100644 index 000000000..a66e385db --- /dev/null +++ b/packages/diffusers-cpu-offloading/benchmarks/v2v_prompt.json @@ -0,0 +1,136 @@ +{ + "subjects": [ + { + "description": "A sleek industrial robot arm with a silver and dark gray metallic finish, featuring multiple articulated joints and a gripper end-effector that holds a clear glass jar filled with water.", + "appearance_details": "The robot arm has visible servo motors at each joint, subtle branding embossed on the upper arm segment, black rubber gaskets at joint connections, and a precision two-finger parallel gripper clamping the jar securely.", + "relationship": "Primary actor performing the pouring action; positioned above and beside the cup.", + "location": "Center-right of frame", + "relative_size": "Large within frame", + "orientation": "Angled toward the left side of frame, gripper tilted downward toward the cup", + "pose": "Extended arm with elbow joint bent, wrist rotated to tilt the jar at a pouring angle", + "action": "Pouring water from a glass jar into a cup", + "state_changes": "Begins in a neutral upright hold, then gradually tilts the jar to pour, and returns slightly upright as the pour finishes.", + "clothing": "", + "expression": "", + "gender": "", + "age": "", + "skin_tone_and_texture": "", + "facial_features": "", + "number_of_subjects": 1, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "A clear glass jar approximately three-quarters full of water, with a wide mouth opening, held by the robot arm's gripper.", + "appearance_details": "Transparent glass with slight green tint, smooth cylindrical body, no label, water visible sloshing gently inside as the jar tilts.", + "relationship": "Held by the robot arm; source of the water being poured.", + "location": "Center-right, elevated above the cup", + "relative_size": "Medium within frame", + "orientation": "Tilting progressively toward the left as the pour occurs", + "pose": "Held at an angle by the gripper", + "action": "Being tilted to pour water", + "state_changes": "Water level decreases as the pour progresses; jar tilts from near-vertical to approximately 45 degrees.", + "clothing": "", + "expression": "", + "gender": "", + "age": "", + "skin_tone_and_texture": "", + "facial_features": "", + "number_of_subjects": 1, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "A white ceramic cup sitting on a flat surface, positioned to receive the poured water.", + "appearance_details": "Simple cylindrical mug shape with a small handle on the right side, matte white glaze, clean and empty at the start.", + "relationship": "Receiving vessel for the water being poured from the jar.", + "location": "Center-left foreground, on the table surface", + "relative_size": "Small within frame", + "orientation": "Upright, opening facing upward", + "pose": "Stationary on the table", + "action": "Receiving water", + "state_changes": "Gradually fills with water as the pour continues.", + "clothing": "", + "expression": "", + "gender": "", + "age": "", + "skin_tone_and_texture": "", + "facial_features": "", + "number_of_subjects": 1, + "number_of_arms": 0, + "number_of_legs": 0 + } + ], + "background_setting": "A clean, minimalist laboratory or workshop environment with a plain light gray wall in the background and a smooth white tabletop surface. The space is uncluttered, with no other objects or distractions visible, contributing to a very quiet and sterile atmosphere. The edges of the table are just visible at the bottom of the frame.", + "lighting": { + "conditions": "Soft, even studio lighting with minimal harsh highlights, creating a clean and controlled look.", + "direction": "Front-lit with slight top-down component, providing even illumination across the robot arm and objects.", + "shadows": "Soft, diffused shadows beneath the cup and the robot arm's base, with gentle shadow on the table surface.", + "illumination_effect": "The lighting emphasizes the metallic sheen of the robot arm and the transparency of the water and glass jar, giving the scene a polished, technical demonstration feel." + }, + "aesthetics": { + "composition": "The robot arm dominates the right half of the frame while the cup sits in the left-center foreground, creating a clear visual flow from right to left following the pouring action. Negative space above and behind keeps focus on the action.", + "color_scheme": "Neutral palette dominated by silver, gray, and white with the clear blue-tinted transparency of water providing subtle color contrast.", + "mood_atmosphere": "Calm, precise, clinical, quietly impressive", + "patterns": "" + }, + "cinematography": { + "camera_motion": "Static", + "framing": "Medium shot", + "camera_angle": "Eye-level, slightly elevated", + "depth_of_field": "Shallow", + "focus": "Sharp focus on the pouring point where water exits the jar and enters the cup", + "lens_focal_length": "50mm equivalent" + }, + "style_medium": "Live-action video", + "artistic_style": "Realistic, clean technical demonstration", + "context": "A robotics demonstration showcasing a robot arm's precision and dexterity in performing a delicate pouring task.", + "actions": [ + { + "time": "0:00-0:02", + "description": "The robot arm holds the jar upright in a steady position above the cup, making small preparatory adjustments to align the jar's mouth with the cup opening." + }, + { + "time": "0:02-0:05", + "description": "The robot arm smoothly tilts the jar, and a steady stream of clear water flows from the jar into the white ceramic cup below. The water stream is smooth and controlled." + }, + { + "time": "0:05-0:07", + "description": "The robot arm gradually returns the jar to a more upright position, the water stream thins and stops, and the cup is now partially filled with water. The arm holds still in its final position." + } + ], + "text_and_signage_elements": [], + "segments": [ + { + "segment_index": 0, + "time_range": "0:00-0:02", + "description": "Opening shot establishes the robot arm holding the jar above the cup. The arm makes slight positional adjustments.", + "key_changes": "Minor wrist rotation as the arm aligns the jar with the cup.", + "camera": "Static, medium shot at eye level." + }, + { + "segment_index": 1, + "time_range": "0:02-0:05", + "description": "The main pouring action occurs as the robot arm tilts the jar and water flows in a controlled stream into the cup.", + "key_changes": "Jar tilts from near-vertical to approximately 45 degrees; water stream begins and maintains a steady flow; cup fills progressively.", + "camera": "Static, maintaining focus on the pouring point." + }, + { + "segment_index": 2, + "time_range": "0:05-0:07", + "description": "The pour concludes as the arm returns the jar upright. The water stream tapers off and the scene settles into stillness.", + "key_changes": "Jar returns toward vertical; water stream narrows and ceases; cup now holds water; arm stabilizes.", + "camera": "Static, same framing held throughout." + } + ], + "transitions": [], + "temporal_caption": "The video opens with a silver robotic arm holding a clear glass jar of water above a white ceramic cup on a clean white table against a gray background. During the first two seconds, the arm makes subtle alignment adjustments. At around two seconds, the wrist joint rotates and the jar begins to tilt, releasing a smooth, steady stream of clear water that arcs downward into the cup. The pouring continues for about three seconds, with the water level in the jar visibly decreasing and the cup gradually filling. Around the five-second mark, the arm begins to level the jar back upright, the stream of water thins to a trickle, then stops entirely. For the remaining two seconds, the robot arm holds the jar still in a slightly tilted resting position, and the filled cup sits motionless on the table.", + "audio_description": "The scene is very quiet. The dominant sound is the gentle splashing and trickling of water as it pours from the jar into the cup, starting softly and becoming slightly louder as the cup fills. There is a faint mechanical whir from the robot arm's servo motors during movement. No music, no speech, and minimal ambient noise, emphasizing the tranquility of the environment.", + "resolution": { + "H": 720, + "W": 1280 + }, + "aspect_ratio": "16,9", + "duration": "7s", + "fps": 24 + } diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py new file mode 100644 index 000000000..1141e9d57 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py @@ -0,0 +1,26 @@ +"""Cosmos 3 Nano streaming execution for Diffusers modular pipelines.""" + +from .config import StreamingConfig +from .metrics import RuntimeStats + +__all__ = ["RuntimeStats", "StreamingConfig"] +__version__ = "0.1.0" + + +def build_pipeline(*args, **kwargs): + """Lazily import the unified Diffusers-dependent pipeline builder.""" + + from .modular import build_pipeline as _build + + return _build(*args, **kwargs) + + +def build_text_to_video_pipeline(*args, **kwargs): + """Compatibility alias for the unified pipeline builder.""" + + from .modular import build_text_to_video_pipeline as _build + + return _build(*args, **kwargs) + + +__all__.extend(["build_pipeline", "build_text_to_video_pipeline"]) diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py new file mode 100644 index 000000000..57e1bf81d --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py @@ -0,0 +1,9 @@ +from .base import CosmosAttentionBackend +from .loader import load_attention_backend +from .sdpa import TorchSdpaBackend + +__all__ = [ + "CosmosAttentionBackend", + "TorchSdpaBackend", + "load_attention_backend", +] diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py new file mode 100644 index 000000000..5add4421e --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +import torch + + +@runtime_checkable +class CosmosAttentionBackend(Protocol): + """Attention contract used by the split Cosmos executor. + + Inputs and outputs use token-major NHD layout: ``[tokens, heads, dim]``. + The backend must support GQA, where K/V have fewer heads than Q. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + is_causal: bool, + ) -> torch.Tensor: + """Return attention output in the same NHD layout as ``q``.""" diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py new file mode 100644 index 000000000..2f4fffd13 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Any + +from .base import CosmosAttentionBackend + + +def load_attention_backend(spec: str, **kwargs: Any) -> CosmosAttentionBackend: + """Instantiate an attention backend from ``module.path:attribute``.""" + + try: + module_name, attribute_name = spec.rsplit(":", 1) + except ValueError as exc: + raise ValueError(f"invalid attention backend {spec!r}; expected 'module.path:ClassName'") from exc + + module = import_module(module_name) + factory = getattr(module, attribute_name) + backend = factory(**kwargs) + if not isinstance(backend, CosmosAttentionBackend): + raise TypeError(f"{spec!r} did not produce a CosmosAttentionBackend-compatible object") + return backend diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py new file mode 100644 index 000000000..b4c19f855 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +class TorchSdpaBackend: + """BF16 reference backend using PyTorch scaled dot-product attention.""" + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + is_causal: bool, + ) -> torch.Tensor: + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("q, k, and v must use [tokens, heads, dim] layout") + if k.shape != v.shape: + raise ValueError(f"k and v shapes must match, got {tuple(k.shape)} and {tuple(v.shape)}") + if q.shape[-1] != k.shape[-1]: + raise ValueError("q and k head dimensions must match") + if q.shape[1] % k.shape[1] != 0: + raise ValueError("query head count must be divisible by KV head count") + + # torch SDPA consumes [batch, heads, tokens, dim]. + q_bhsd = q.transpose(0, 1).unsqueeze(0) + k_bhsd = k.transpose(0, 1).unsqueeze(0) + v_bhsd = v.transpose(0, 1).unsqueeze(0) + + try: + output = F.scaled_dot_product_attention( + q_bhsd, + k_bhsd, + v_bhsd, + is_causal=is_causal, + enable_gqa=q.shape[1] != k.shape[1], + ) + except TypeError: + # Compatibility fallback for PyTorch releases without enable_gqa. + groups = q.shape[1] // k.shape[1] + if groups != 1: + k_bhsd = k_bhsd.repeat_interleave(groups, dim=1) + v_bhsd = v_bhsd.repeat_interleave(groups, dim=1) + output = F.scaled_dot_product_attention( + q_bhsd, + k_bhsd, + v_bhsd, + is_causal=is_causal, + ) + + return output.squeeze(0).transpose(0, 1).contiguous() diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py new file mode 100644 index 000000000..31934a40c --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence + +import torch +from torch import nn + +from .attention.base import CosmosAttentionBackend + +type RotaryPair = tuple[torch.Tensor, torch.Tensor] +type PrefixKV = tuple[torch.Tensor, torch.Tensor] +MEMORY_SERIAL_MLP_CHUNK_TOKENS = 8_192 + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def apply_rotary(x: torch.Tensor, rotary: RotaryPair) -> torch.Tensor: + cos, sin = rotary + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return x * cos + _rotate_half(x) * sin + + +class UnderstandingBlockView(nn.Module): + """Understanding-only execution view over one original Cosmos decoder layer.""" + + def __init__(self, layer: nn.Module, backend: CosmosAttentionBackend) -> None: + super().__init__() + attention = layer.self_attn + + # Register only understanding-side modules. Registering the complete + # attention module would make its generation projections part of this + # offload group as well. + self.input_layernorm = layer.input_layernorm + self.post_attention_layernorm = layer.post_attention_layernorm + self.mlp = layer.mlp + self.to_q = attention.to_q + self.to_k = attention.to_k + self.to_v = attention.to_v + self.to_out = attention.to_out + self.norm_q = attention.norm_q + self.norm_k = attention.norm_k + self.k_norm_und_for_gen = attention.k_norm_und_for_gen + + self.num_attention_heads = attention.num_attention_heads + self.num_key_value_heads = attention.num_key_value_heads + self.head_dim = attention.head_dim + self._backend = backend + + def _one_branch(self, hidden: torch.Tensor, rotary: RotaryPair) -> tuple[torch.Tensor, PrefixKV]: + normalized = self.input_layernorm(hidden) + q = self.to_q(normalized).view(-1, self.num_attention_heads, self.head_dim) + k = self.to_k(normalized).view(-1, self.num_key_value_heads, self.head_dim) + v = self.to_v(normalized).view(-1, self.num_key_value_heads, self.head_dim) + + q = self.norm_q(q) + k = self.norm_k(k) + k_for_generation = self.k_norm_und_for_gen(k) if self.k_norm_und_for_gen is not None else k + + q = apply_rotary(q, rotary) + k = apply_rotary(k, rotary) + k_for_generation = apply_rotary(k_for_generation, rotary) + + attention_output = self._backend(q, k, v, is_causal=True).flatten(-2, -1) + residual = hidden + self.to_out(attention_output) + output = residual + self.mlp(self.post_attention_layernorm(residual)) + return output, (k_for_generation, v) + + def forward( + self, + conditional_hidden: torch.Tensor, + conditional_rotary: RotaryPair, + unconditional_hidden: torch.Tensor | None = None, + unconditional_rotary: RotaryPair | None = None, + ) -> tuple[ + torch.Tensor, + PrefixKV, + torch.Tensor | None, + PrefixKV | None, + ]: + conditional_hidden, conditional_prefix = self._one_branch(conditional_hidden, conditional_rotary) + + unconditional_prefix = None + if unconditional_hidden is not None: + if unconditional_rotary is None: + raise ValueError("unconditional_rotary is required with unconditional_hidden") + unconditional_hidden, unconditional_prefix = self._one_branch(unconditional_hidden, unconditional_rotary) + + return conditional_hidden, conditional_prefix, unconditional_hidden, unconditional_prefix + + +class GenerationBlockView(nn.Module): + """Generation-only execution view over one original Cosmos decoder layer.""" + + def __init__(self, layer: nn.Module, backend: CosmosAttentionBackend) -> None: + super().__init__() + attention = layer.self_attn + + self.input_layernorm = layer.input_layernorm_moe_gen + self.post_attention_layernorm = layer.post_attention_layernorm_moe_gen + self.mlp = layer.mlp_moe_gen + self.to_q = attention.add_q_proj + self.to_k = attention.add_k_proj + self.to_v = attention.add_v_proj + self.to_out = attention.to_add_out + self.norm_q = attention.norm_added_q + self.norm_k = attention.norm_added_k + + self.num_attention_heads = attention.num_attention_heads + self.num_key_value_heads = attention.num_key_value_heads + self.head_dim = attention.head_dim + self._backend = backend + + def _one_branch( + self, + hidden: torch.Tensor, + rotary: RotaryPair, + understanding_prefix: PrefixKV, + *, + mlp_chunk_size: int | None = None, + ) -> torch.Tensor: + normalized = self.input_layernorm(hidden) + q = self.to_q(normalized).view(-1, self.num_attention_heads, self.head_dim) + k = self.to_k(normalized).view(-1, self.num_key_value_heads, self.head_dim) + v = self.to_v(normalized).view(-1, self.num_key_value_heads, self.head_dim) + + q = apply_rotary(self.norm_q(q), rotary) + k = apply_rotary(self.norm_k(k), rotary) + + prefix_k, prefix_v = understanding_prefix + all_k = torch.cat((prefix_k, k), dim=0) + all_v = torch.cat((prefix_v, v), dim=0) + attention_output = self._backend(q, all_k, all_v, is_causal=False).flatten(-2, -1) + + residual = hidden + self.to_out(attention_output) + normalized = self.post_attention_layernorm(residual) + if mlp_chunk_size is None or normalized.shape[0] <= mlp_chunk_size: + return residual + self.mlp(normalized) + + # The unchunked 720p Transfer MLP needs a 1.31 GiB contiguous + # intermediate. Process token rows independently and accumulate into + # the existing residual allocation. Linear and pointwise MLP + # operations do not mix information across the token dimension. + for start in range(0, normalized.shape[0], mlp_chunk_size): + end = min(start + mlp_chunk_size, normalized.shape[0]) + residual[start:end].add_(self.mlp(normalized[start:end])) + return residual + + def forward( + self, + conditional_hidden: torch.Tensor, + conditional_rotary: RotaryPair, + conditional_prefix: PrefixKV, + unconditional_hidden: torch.Tensor | None = None, + unconditional_rotary: RotaryPair | None = None, + unconditional_prefix: PrefixKV | None = None, + extra_hiddens: list[torch.Tensor] | None = None, + extra_rotaries: list[RotaryPair] | None = None, + extra_prefixes: list[PrefixKV] | None = None, + stage_branches: bool = False, + execution_device: torch.device | None = None, + offload_device: torch.device | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None] | tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + def run_branch(hidden: torch.Tensor, rotary: RotaryPair, prefix: PrefixKV) -> torch.Tensor: + if not stage_branches: + return self._one_branch(hidden, rotary, prefix) + if execution_device is None or offload_device is None: + raise ValueError("execution_device and offload_device are required when staging branches") + + device_hidden = hidden.to( + execution_device, + non_blocking=hidden.device.type == "cpu" and hidden.is_pinned(), + ) + device_output = self._one_branch( + device_hidden, + rotary, + prefix, + mlp_chunk_size=MEMORY_SERIAL_MLP_CHUNK_TOKENS, + ) + if ( + hidden.device == offload_device + and hidden.shape == device_output.shape + and hidden.dtype == device_output.dtype + ): + hidden.copy_(device_output, non_blocking=hidden.device.type == "cpu" and hidden.is_pinned()) + return hidden + return device_output.to(offload_device) + + conditional_hidden = run_branch(conditional_hidden, conditional_rotary, conditional_prefix) + + if unconditional_hidden is not None: + if unconditional_rotary is None or unconditional_prefix is None: + raise ValueError("unconditional rotary and prefix are required with unconditional_hidden") + unconditional_hidden = run_branch( + unconditional_hidden, + unconditional_rotary, + unconditional_prefix, + ) + if extra_hiddens is not None: + if extra_rotaries is None or extra_prefixes is None: + raise ValueError("extra rotary embeddings and prefixes are required with extra branches") + if not (len(extra_hiddens) == len(extra_rotaries) == len(extra_prefixes)): + raise ValueError("extra generation branch inputs must have the same length") + extra_hiddens = [ + run_branch(hidden, rotary, prefix) + for hidden, rotary, prefix in zip(extra_hiddens, extra_rotaries, extra_prefixes) + ] + return conditional_hidden, unconditional_hidden, extra_hiddens + return conditional_hidden, unconditional_hidden + + +class UnderstandingStack(nn.Module): + """Runs both CFG understanding branches once and returns per-layer K/V.""" + + def __init__( + self, + layers: Sequence[nn.Module], + backend: CosmosAttentionBackend, + block_callback: Callable[[str, int], None] | None = None, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList(UnderstandingBlockView(layer, backend) for layer in layers) + self._block_callback = block_callback + # Diffusers streamed group offload needs an unmatched root group to own + # the lazy-prefetch trace. Identity is intentionally parameter-free. + self._prefetch_anchor = nn.Identity() + + def forward( + self, + conditional_hidden: torch.Tensor, + conditional_rotary: RotaryPair, + unconditional_hidden: torch.Tensor | None = None, + unconditional_rotary: RotaryPair | None = None, + ) -> tuple[list[PrefixKV], list[PrefixKV] | None]: + conditional_cache: list[PrefixKV] = [] + unconditional_cache: list[PrefixKV] | None = [] if unconditional_hidden is not None else None + + for index, block in enumerate(self.blocks): + conditional_hidden, conditional_prefix, unconditional_hidden, unconditional_prefix = block( + conditional_hidden, + conditional_rotary, + unconditional_hidden, + unconditional_rotary, + ) + conditional_cache.append(conditional_prefix) + if self._block_callback is not None: + self._block_callback("understanding", index) + if unconditional_cache is not None: + if unconditional_prefix is None: + raise RuntimeError("understanding block did not return an unconditional prefix") + unconditional_cache.append(unconditional_prefix) + + return conditional_cache, unconditional_cache + + +class GenerationStack(nn.Module): + """Runs both CFG generation branches with one onload per decoder block.""" + + def __init__( + self, + layers: Sequence[nn.Module], + backend: CosmosAttentionBackend, + block_callback: Callable[[str, int], None] | None = None, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList(GenerationBlockView(layer, backend) for layer in layers) + self._block_callback = block_callback + self._prefetch_anchor = nn.Identity() + + def forward( + self, + conditional_hidden: torch.Tensor, + conditional_rotary: RotaryPair, + conditional_cache: Sequence[PrefixKV], + unconditional_hidden: torch.Tensor | None = None, + unconditional_rotary: RotaryPair | None = None, + unconditional_cache: Sequence[PrefixKV] | None = None, + extra_hiddens: list[torch.Tensor] | None = None, + extra_rotaries: list[RotaryPair] | None = None, + extra_caches: list[Sequence[PrefixKV]] | None = None, + stage_branches: bool = False, + execution_device: torch.device | None = None, + offload_device: torch.device | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None] | tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + if len(conditional_cache) != len(self.blocks): + raise ValueError("conditional K/V cache does not match the decoder depth") + if unconditional_hidden is not None and ( + unconditional_cache is None or len(unconditional_cache) != len(self.blocks) + ): + raise ValueError("unconditional K/V cache does not match the decoder depth") + if extra_hiddens is not None: + if extra_rotaries is None or extra_caches is None: + raise ValueError("extra rotary embeddings and K/V caches are required with extra branches") + if not (len(extra_hiddens) == len(extra_rotaries) == len(extra_caches)): + raise ValueError("extra generation branch inputs must have the same length") + if any(len(cache) != len(self.blocks) for cache in extra_caches): + raise ValueError("an extra K/V cache does not match the decoder depth") + + for index, block in enumerate(self.blocks): + unconditional_prefix = None if unconditional_cache is None else unconditional_cache[index] + if extra_hiddens is None: + conditional_hidden, unconditional_hidden = block( + conditional_hidden, + conditional_rotary, + conditional_cache[index], + unconditional_hidden, + unconditional_rotary, + unconditional_prefix, + stage_branches=stage_branches, + execution_device=execution_device, + offload_device=offload_device, + ) + else: + conditional_hidden, unconditional_hidden, extra_hiddens = block( + conditional_hidden, + conditional_rotary, + conditional_cache[index], + unconditional_hidden, + unconditional_rotary, + unconditional_prefix, + extra_hiddens, + extra_rotaries, + [cache[index] for cache in extra_caches], + stage_branches=stage_branches, + execution_device=execution_device, + offload_device=offload_device, + ) + if self._block_callback is not None: + self._block_callback("generation", index) + if extra_hiddens is not None: + return conditional_hidden, unconditional_hidden, extra_hiddens + return conditional_hidden, unconditional_hidden diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py new file mode 100644 index 000000000..4d0ae438f --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + + +@dataclass(frozen=True, slots=True) +class StreamingConfig: + """Runtime policy for Cosmos 3 split execution. + + The model weights remain in their checkpoint dtype. ``device`` and + ``offload_device`` control placement only; this class does not enable model + weight quantization. + """ + + device: str = "cuda" + offload_device: str = "cpu" + enable_group_offload: bool = True + use_stream: bool = True + record_stream: bool = False + low_cpu_mem_usage: bool = True + vae_leaf_offload: bool = True + memory_profile: Literal["auto", "performance", "low_memory"] = "auto" + attention_backend: str = "cosmos3_streaming.attention.sdpa:TorchSdpaBackend" + attention_backend_kwargs: dict[str, Any] = field(default_factory=dict) + empty_cuda_cache_on_end: bool = False + memory_trace_path: str | None = None + memory_trace_blocks: bool = False + memory_trace_synchronize: bool = False + + def __post_init__(self) -> None: + if not self.device: + raise ValueError("device must be non-empty") + if self.offload_device != "cpu": + raise ValueError("the first implementation only supports CPU weight storage") + if self.record_stream and not self.use_stream: + raise ValueError("record_stream=True requires use_stream=True") + if ":" not in self.attention_backend: + raise ValueError("attention_backend must use 'module.path:ClassName' syntax") + if self.memory_trace_blocks and self.memory_trace_path is None: + raise ValueError("memory_trace_blocks=True requires memory_trace_path") + if self.memory_profile not in ("auto", "performance", "low_memory"): + raise ValueError(f"unsupported memory profile: {self.memory_profile}") diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py new file mode 100644 index 000000000..2da4efe9f --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +DEFAULT_FP8_MODEL_ID = "nvidia/Cosmos3-Nano" +DEFAULT_FP8_REVISION = "fp8" + + +@dataclass(frozen=True, slots=True) +class FP8LoadReport: + """Evidence that ModelOpt restored a usable pre-quantized transformer.""" + + fp8_parameters: int + quantizer_modules: int + meta_tensors: int + + +def materialize_fp8_checkpoint( + model_id: str, + *, + revision: str | None, + local_files_only: bool, +) -> Path: + """Resolve an FP8 checkpoint to a local directory. + + ModelOpt 0.44 restores ``modelopt_state.pth`` relative to the model path and + does not reliably honor the modular loader's separate ``subfolder`` value. + Materializing the snapshot lets the adapter load ``transformer/`` through + its complete path while keeping all other components on the normal + Diffusers modular path. + """ + + local_path = Path(model_id).expanduser() + if local_path.is_dir(): + checkpoint_dir = local_path.resolve() + else: + from huggingface_hub import snapshot_download + + # Nano publishes BF16 on its default branch and FP8 on a separate one. + # Do not override a caller's pinned revision or another repository's default. + if revision is None and model_id == DEFAULT_FP8_MODEL_ID: + revision = DEFAULT_FP8_REVISION + checkpoint_dir = Path( + snapshot_download( + model_id, + revision=revision, + local_files_only=local_files_only, + ) + ).resolve() + + modelopt_state = checkpoint_dir / "transformer" / "modelopt_state.pth" + if not modelopt_state.is_file(): + raise RuntimeError( + f"FP8 was requested, but the checkpoint has no ModelOpt state at {modelopt_state}. " + f"Use {DEFAULT_FP8_MODEL_ID} with revision={DEFAULT_FP8_REVISION!r}, " + "or a local Diffusers FP8 snapshot." + ) + return checkpoint_dir + + +def activate_modelopt_fp8() -> None: + """Register ModelOpt's Hugging Face restore path before model loading.""" + + try: + # Importing the backend registers the FP8 GEMM implementation used by + # the restored quantized Linear modules. + import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401 + from modelopt.torch.opt import enable_huggingface_checkpointing + except ImportError as error: + raise ImportError( + "FP8 checkpoint loading requires NVIDIA ModelOpt. In your active Diffusers environment, run " + "`uv pip install nvidia-modelopt==0.44.0` (or `python -m pip install nvidia-modelopt==0.44.0`). " + "Alternatively install this package with its [fp8] extra." + ) from error + + enable_huggingface_checkpointing() + + +def load_fp8_transformer( + checkpoint_dir: Path, + *, + dtype: torch.dtype, +) -> torch.nn.Module: + """Restore the ModelOpt transformer from its complete local component path.""" + + activate_modelopt_fp8() + + from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer + + transformer = Cosmos3OmniTransformer.from_pretrained( + str(checkpoint_dir / "transformer"), + dtype=dtype, + local_files_only=True, + ) + verify_fp8_transformer(transformer) + return transformer + + +def verify_fp8_transformer(transformer: torch.nn.Module) -> FP8LoadReport: + """Reject partial restores that contain FP8 tensors but no quantizer graph.""" + + named_tensors: list[tuple[str, Any]] = list(transformer.named_parameters()) + list(transformer.named_buffers()) + fp8_parameters = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters()) + quantizer_modules = sum("quantizer" in name.lower() for name, _ in transformer.named_modules()) + meta_tensors = sum(tensor.is_meta for _, tensor in named_tensors) + report = FP8LoadReport( + fp8_parameters=fp8_parameters, + quantizer_modules=quantizer_modules, + meta_tensors=meta_tensors, + ) + if not fp8_parameters or not quantizer_modules or meta_tensors: + raise RuntimeError( + "invalid Cosmos 3 FP8 restore: " + f"fp8_parameters={fp8_parameters}, quantizer_modules={quantizer_modules}, " + f"meta_tensors={meta_tensors}" + ) + return report diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py new file mode 100644 index 000000000..55a323ab5 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +import torch + + +class CudaMemoryTrace: + """Interval and cumulative CUDA-memory telemetry for one pipeline request. + + Each checkpoint closes the interval that began at the preceding checkpoint. + PyTorch's allocator peak is reset after the sample, so ``interval_peak_*`` + describes that interval rather than the whole process. ``device_used_*`` + comes from the CUDA driver and therefore includes CUDA context and non-PyTorch + allocations as well as this process's allocator state. + """ + + def __init__(self, path: str | Path, device: torch.device, *, synchronize: bool = False) -> None: + self.path = Path(path) + self.device = device + self.synchronize = synchronize + self._started_at = 0.0 + self._sequence = 0 + self._global_peak_allocated = 0 + self._global_peak_reserved = 0 + self._global_peak_device_used = 0 + self._active = False + + def begin(self) -> None: + self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.write_text("") + self._started_at = time.perf_counter() + self._sequence = 0 + self._global_peak_allocated = 0 + self._global_peak_reserved = 0 + self._global_peak_device_used = 0 + self._active = True + if self.device.type == "cuda": + with torch.cuda.device(self.device): + torch.cuda.reset_peak_memory_stats() + self.checkpoint("request.begin") + + def checkpoint(self, event: str, **fields: Any) -> None: + if not self._active: + return + + sample: dict[str, Any] = { + "sequence": self._sequence, + "elapsed_seconds": time.perf_counter() - self._started_at, + "event": event, + **fields, + } + if self.device.type == "cuda": + if self.synchronize: + torch.cuda.synchronize(self.device) + with torch.cuda.device(self.device): + allocated = torch.cuda.memory_allocated() + reserved = torch.cuda.memory_reserved() + interval_peak_allocated = torch.cuda.max_memory_allocated() + interval_peak_reserved = torch.cuda.max_memory_reserved() + free, total = torch.cuda.mem_get_info() + device_used = total - free + torch.cuda.reset_peak_memory_stats() + + self._global_peak_allocated = max(self._global_peak_allocated, interval_peak_allocated) + self._global_peak_reserved = max(self._global_peak_reserved, interval_peak_reserved) + self._global_peak_device_used = max(self._global_peak_device_used, device_used) + sample.update( + { + "allocated_bytes": allocated, + "reserved_bytes": reserved, + "interval_peak_allocated_bytes": interval_peak_allocated, + "interval_peak_reserved_bytes": interval_peak_reserved, + "device_used_bytes": device_used, + "device_total_bytes": total, + "global_peak_allocated_bytes": self._global_peak_allocated, + "global_peak_reserved_bytes": self._global_peak_reserved, + "global_peak_device_used_bytes": self._global_peak_device_used, + } + ) + + with self.path.open("a") as handle: + handle.write(json.dumps(sample, sort_keys=True) + "\n") + self._sequence += 1 + + def end(self) -> None: + if not self._active: + return + self.checkpoint("request.end") + self._active = False diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py new file mode 100644 index 000000000..be4c3fee2 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(slots=True) +class RuntimeStats: + """Low-overhead counters for one completed or active request. + + Byte counters cover decoder-block parameters only. CUDA/Nsight profiling + remains the source of truth for physical PCIe traffic. + """ + + understanding_prefills: int = 0 + denoising_calls: int = 0 + understanding_weight_bytes: int = 0 + generation_weight_bytes: int = 0 + estimated_block_h2d_bytes: int = 0 + + @property + def estimated_block_h2d_gib(self) -> float: + return self.estimated_block_h2d_bytes / 2**30 diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py new file mode 100644 index 000000000..990890e50 --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py @@ -0,0 +1,656 @@ +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +from diffusers import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline +from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer +from diffusers.modular_pipelines import BlockState, ComponentSpec, InputParam, ModularPipelineBlocks, OutputParam + +from .config import StreamingConfig +from .fp8 import DEFAULT_FP8_MODEL_ID, load_fp8_transformer, materialize_fp8_checkpoint +from .offloading import offload_groups +from .runtime import StreamingCosmosRuntime + + +class StreamingCosmos3LoopDenoiser(ModularPipelineBlocks): + """External streamed replacement for Diffusers' multimodal denoiser.""" + + model_name = "cosmos3-omni" + + def __init__(self, runtime: StreamingCosmosRuntime) -> None: + super().__init__() + self.runtime = runtime + + @property + def description(self) -> str: + return "Predicts Cosmos 3 modality velocities with cached understanding and streamed generation blocks." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("denoiser_input_fields"), + InputParam( + name="guidance_scale", + type_hint=float, + default=6.0, + description="Scale for classifier-free guidance.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "velocity_vision", + type_hint=torch.Tensor, + description="Predicted velocity for vision latents.", + ), + OutputParam("velocity_sound", type_hint=torch.Tensor, description="Predicted velocity for sound latents."), + OutputParam( + "velocity_action", type_hint=torch.Tensor, description="Predicted velocity for action latents." + ), + ] + + @torch.no_grad() + def __call__( + self, + components: Cosmos3OmniModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ) -> tuple[Cosmos3OmniModularPipeline, BlockState]: + del t # The per-step timestep tensor is already present in block_state. + self.runtime.trace_checkpoint("denoise_step.start", step=i) + + denoiser_input_fields = block_state.denoiser_input_fields + loop_input_fields = block_state.as_dict() + has_sound = "sound_tokens" in loop_input_fields + has_action = "action_tokens" in loop_input_fields + + transformer_args = set(inspect.signature(components.transformer.forward).parameters) + do_cfg = block_state.guidance_scale != 1.0 + pass_names = ["cond", "uncond"] if do_cfg else ["cond"] + branches = { + pass_name: self._collect_transformer_kwargs( + pass_name, + denoiser_input_fields, + loop_input_fields, + transformer_args, + ) + for pass_name in pass_names + } + predictions = self.runtime.predict_modalities(components.transformer, branches) + + velocities: dict[str, tuple[torch.Tensor, torch.Tensor | None, None]] = {} + for pass_name, prediction in predictions.items(): + velocities[pass_name] = components._mask_velocity_predictions( + prediction.vision, + prediction.sound, + vision_condition_mask=[loop_input_fields["vision_condition_mask"]], + sound_condition_mask=[loop_input_fields["sound_condition_mask"]] if has_sound else None, + preds_action=prediction.action, + action_condition_mask=[loop_input_fields["action_condition_mask"]] if has_action else None, + raw_action_dim=loop_input_fields.get("raw_action_dim_resolved"), + ) + + cond_velocity, cond_velocity_sound, cond_velocity_action = velocities["cond"] + if do_cfg: + uncond_velocity, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"] + block_state.velocity_vision = uncond_velocity + block_state.guidance_scale * ( + cond_velocity - uncond_velocity + ) + block_state.velocity_sound = ( + uncond_velocity_sound + block_state.guidance_scale * (cond_velocity_sound - uncond_velocity_sound) + if has_sound + else None + ) + block_state.velocity_action = ( + uncond_velocity_action + block_state.guidance_scale * (cond_velocity_action - uncond_velocity_action) + if has_action + else None + ) + else: + block_state.velocity_vision = cond_velocity + block_state.velocity_sound = cond_velocity_sound if has_sound else None + block_state.velocity_action = cond_velocity_action if has_action else None + self.runtime.trace_checkpoint("denoise_step.end", step=i) + return components, block_state + + @staticmethod + def _collect_transformer_kwargs( + pass_name: str, + denoiser_input_fields: dict[str, Any], + loop_input_fields: dict[str, Any], + transformer_args: set[str], + ) -> dict[str, Any]: + kwargs: dict[str, Any] = {} + prefix = f"{pass_name}_" + for field_name, field_value in denoiser_input_fields.items(): + if field_name.startswith(prefix): + transformer_field_name = field_name.removeprefix(prefix) + if transformer_field_name.endswith("_segment"): + kwargs.update(field_value) + else: + kwargs[transformer_field_name] = field_value + elif field_name in transformer_args: + kwargs[field_name] = field_value + + kwargs.update( + { + field_name: field_value + for field_name, field_value in loop_input_fields.items() + if field_name in transformer_args + } + ) + kwargs = {name: value for name, value in kwargs.items() if name in transformer_args} + + required = { + "input_ids", + "text_indexes", + "position_ids", + "und_len", + "sequence_length", + "vision_tokens", + "vision_token_shapes", + "vision_sequence_indexes", + "vision_mse_loss_indexes", + "vision_timesteps", + "vision_noisy_frame_indexes", + } + missing = sorted(required - kwargs.keys()) + if missing: + raise KeyError(f"missing Cosmos transformer inputs for {pass_name}: {missing}") + return kwargs + + +class StreamingCosmos3TransferLoopDenoiser(ModularPipelineBlocks): + """Streamed replacement for the three-axis transfer denoiser.""" + + model_name = "cosmos3-omni" + + def __init__(self, runtime: StreamingCosmosRuntime) -> None: + super().__init__() + self.runtime = runtime + + @property + def description(self) -> str: + return "Predicts transfer velocity with streamed control and text guidance branches." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("denoiser_input_fields"), + InputParam(name="vision_tokens_full", type_hint=list[torch.Tensor], required=True), + InputParam(name="vision_tokens_target", type_hint=list[torch.Tensor], required=True), + InputParam(name="vision_timesteps", type_hint=torch.Tensor, required=True), + InputParam(name="velocity_mask", type_hint=torch.Tensor, required=True), + InputParam(name="guidance_scale", type_hint=float, default=6.0), + InputParam(name="control_guidance", type_hint=float, default=1.0), + InputParam(name="guidance_interval", type_hint=tuple, default=None), + InputParam(name="control_guidance_interval", type_hint=tuple, default=None), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted transfer velocity.")] + + @staticmethod + def _transformer_kwargs( + static: dict[str, Any], + vision_tokens: list[torch.Tensor], + vision_timesteps: torch.Tensor, + ) -> dict[str, Any]: + return { + "input_ids": static["input_ids"], + "text_indexes": static["text_indexes"], + "position_ids": static["position_ids"], + "und_len": static["und_len"], + "sequence_length": static["sequence_length"], + "vision_tokens": vision_tokens, + "vision_token_shapes": static["vision_token_shapes"], + "vision_sequence_indexes": static["vision_sequence_indexes"], + "vision_mse_loss_indexes": static["vision_mse_loss_indexes"], + "vision_timesteps": vision_timesteps, + "vision_noisy_frame_indexes": static["vision_noisy_frame_indexes"], + } + + @staticmethod + def _active(interval: tuple | None, timestep: torch.Tensor) -> bool: + return interval is None or float(interval[0]) <= float(timestep.item()) <= float(interval[1]) + + @torch.no_grad() + def __call__( + self, + components: Cosmos3OmniModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ) -> tuple[Cosmos3OmniModularPipeline, BlockState]: + self.runtime.trace_checkpoint("transfer_denoise_step.start", step=i) + step_guidance = block_state.guidance_scale if self._active(block_state.guidance_interval, t) else 1.0 + step_control = block_state.control_guidance if self._active(block_state.control_guidance_interval, t) else 1.0 + needs_text_cfg = step_guidance > 1.0 + needs_control_cfg = step_control != 1.0 + + static = block_state.denoiser_input_fields + branches = { + "cond_full": self._transformer_kwargs( + static["cond_full_static"], + block_state.vision_tokens_full, + block_state.vision_timesteps, + ) + } + if needs_control_cfg: + branches["cond_no_control"] = self._transformer_kwargs( + static["cond_no_control_static"], + block_state.vision_tokens_target, + block_state.vision_timesteps, + ) + if needs_text_cfg: + branches["uncond_full"] = self._transformer_kwargs( + static["uncond_full_static"], + block_state.vision_tokens_full, + block_state.vision_timesteps, + ) + predictions = self.runtime.predict_modalities(components.transformer, branches) + cond_full = predictions["cond_full"].vision[-1] + + if needs_control_cfg: + cond_no_control = predictions["cond_no_control"].vision[-1] + if needs_text_cfg: + uncond_full = predictions["uncond_full"].vision[-1] + + if needs_control_cfg and needs_text_cfg: + control_cond = cond_no_control + step_control * (cond_full - cond_no_control) + velocity = uncond_full + step_guidance * (control_cond - uncond_full) + elif needs_control_cfg: + velocity = cond_no_control + step_control * (cond_full - cond_no_control) + elif needs_text_cfg: + velocity = uncond_full + step_guidance * (cond_full - uncond_full) + else: + velocity = cond_full + + block_state.velocity = velocity * block_state.velocity_mask + self.runtime.trace_checkpoint("transfer_denoise_step.end", step=i) + return components, block_state + + +class StreamingCosmos3Pipeline(Cosmos3OmniModularPipeline): + """Thin device/lifecycle adapter over the official modular pipeline.""" + + def __init__( + self, + *args: Any, + streaming_runtime: StreamingCosmosRuntime, + execution_device: str | torch.device = "cuda", + **kwargs: Any, + ) -> None: + self._streaming_runtime = streaming_runtime + self._streaming_execution_device = torch.device(execution_device) + super().__init__(*args, **kwargs) + + @property + def _execution_device(self) -> torch.device: + return self._streaming_execution_device + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self._streaming_runtime.begin_request() + failed = True + try: + result = super().__call__(*args, **kwargs) + failed = False + return result + finally: + try: + if failed: + # An encoder/decoder can fail outside the streamed runtime. + # VAE leaf hooks have no prefetch trace and remain reusable. + if self.vae is not None: + offload_groups(self.vae) + if self.sound_tokenizer is not None: + self.sound_tokenizer.to(self._streaming_runtime.offload_device) + finally: + self._streaming_runtime.end_request(failed=failed) + + def _apply_video_safety_check(self, *args: Any, **kwargs: Any) -> Any: + if self._streaming_runtime.release_cache_before_safety: + self._streaming_runtime.trace_checkpoint("pre_safety_cache_clear.start") + torch.cuda.empty_cache() + self._streaming_runtime.trace_checkpoint("pre_safety_cache_clear.end") + self._streaming_runtime.trace_checkpoint("safety_check.start") + try: + return super()._apply_video_safety_check(*args, **kwargs) + finally: + self._streaming_runtime.trace_checkpoint("safety_check.end") + + +class _MemoryTracedBlock(ModularPipelineBlocks): + """Transparent telemetry wrapper around an official modular block.""" + + def __init__(self, block: ModularPipelineBlocks, runtime: StreamingCosmosRuntime, stage: str) -> None: + super().__init__() + self.block = block + self.runtime = runtime + self.stage = stage + self.model_name = block.model_name + + @property + def description(self) -> str: + return self.block.description + + @property + def expected_components(self) -> list[ComponentSpec]: + return self.block.expected_components + + @property + def expected_configs(self) -> list[Any]: + return self.block.expected_configs + + @property + def inputs(self) -> list[InputParam]: + return self.block.inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return self.block.intermediate_outputs + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.runtime.trace_checkpoint(f"pipeline.{self.stage}.start") + result = self.block(*args, **kwargs) + self.runtime.trace_checkpoint(f"pipeline.{self.stage}.end") + return result + + +@dataclass(frozen=True, slots=True) +class _VaeMemoryPolicy: + tile_size: int | None + tile_stride: int | None + release_cache_before_safety: bool + + +def _resolve_vae_memory_policy(profile: str, total_vram_bytes: int) -> _VaeMemoryPolicy: + """Resolve internal VAE details from one researcher-facing memory profile.""" + + if profile == "auto" and total_vram_bytes < 14 * 2**30: + return _VaeMemoryPolicy(tile_size=512, tile_stride=384, release_cache_before_safety=True) + if profile == "low_memory" or (profile == "auto" and total_vram_bytes < 20 * 2**30): + return _VaeMemoryPolicy(tile_size=768, tile_stride=576, release_cache_before_safety=True) + return _VaeMemoryPolicy(tile_size=None, tile_stride=None, release_cache_before_safety=True) + + +def _vision_denoise_loop(blocks: Cosmos3OmniBlocks) -> ModularPipelineBlocks: + """Return the shared vision-only loop in Diffusers' auto-routing Cosmos graph.""" + + try: + return blocks.sub_blocks["denoise"].sub_blocks["vision"].sub_blocks["denoise"] + except KeyError as error: + raise RuntimeError( + "the pinned Cosmos modular block contract changed: expected " + "'denoise.vision.denoise' in the auto-routing workflow" + ) from error + + +def _vision_sound_denoise_loop(blocks: Cosmos3OmniBlocks) -> ModularPipelineBlocks: + """Return the official joint vision-and-sound loop in the auto-routing graph.""" + + try: + return blocks.sub_blocks["denoise"].sub_blocks["vision_sound"].sub_blocks["denoise"] + except KeyError as error: + raise RuntimeError( + "the pinned Cosmos modular block contract changed: expected " + "'denoise.vision_sound.denoise' in the auto-routing workflow" + ) from error + + +def _vision_action_denoise_loop(blocks: Cosmos3OmniBlocks) -> ModularPipelineBlocks: + """Return the official joint vision-and-action loop in the auto-routing graph.""" + + try: + return blocks.sub_blocks["denoise"].sub_blocks["vision_action"].sub_blocks["denoise"] + except KeyError as error: + raise RuntimeError( + "the pinned Cosmos modular block contract changed: expected " + "'denoise.vision_action.denoise' in the auto-routing workflow" + ) from error + + +def _vision_sound_action_denoise_loop(blocks: Cosmos3OmniBlocks) -> ModularPipelineBlocks: + """Return the joint vision, sound, and action loop in the auto-routing graph.""" + + try: + return blocks.sub_blocks["denoise"].sub_blocks["vision_sound_action"].sub_blocks["denoise"] + except KeyError as error: + raise RuntimeError( + "the pinned Cosmos modular block contract changed: expected " + "'denoise.vision_sound_action.denoise' in the auto-routing workflow" + ) from error + + +def _transfer_denoise_loop(blocks: Cosmos3OmniBlocks) -> ModularPipelineBlocks: + """Return the per-chunk transfer loop in Diffusers' auto-routing graph.""" + + try: + return blocks.sub_blocks["denoise"].sub_blocks["transfer"].sub_blocks["chunk_denoise"].sub_blocks["denoise"] + except KeyError as error: + raise RuntimeError( + "the pinned Cosmos modular block contract changed: expected " + "'denoise.transfer.chunk_denoise.denoise' in the auto-routing workflow" + ) from error + + +def _fp8_component_load_overrides( + pipe: Cosmos3OmniModularPipeline, + checkpoint_dir: str, +) -> tuple[dict[str, str], dict[str, None]]: + """Use complete local paths for non-transformer FP8 components. + + The official FP8 snapshot also contains a root ``modelopt_state.pth`` for + its non-Diffusers loading path. ModelOpt 0.44 ignores ``subfolder=`` while + looking for that file, so passing ``root + subfolder`` would incorrectly + try to restore the transformer quantizer graph into the VAE and tokenizer. + """ + + root = Path(checkpoint_dir) + pretrained_paths: dict[str, str] = {} + subfolders: dict[str, None] = {} + for name, spec in pipe._component_specs.items(): + if name == "transformer" or spec.default_creation_method != "from_pretrained": + continue + if not spec.subfolder: + continue + component_dir = root / spec.subfolder + if component_dir.is_dir(): + pretrained_paths[name] = str(component_dir) + subfolders[name] = None + return pretrained_paths, subfolders + + +def build_pipeline( + model_id: str | None = None, + *, + config: StreamingConfig | None = None, + dtype: torch.dtype = torch.bfloat16, + checkpoint_precision: Literal["bf16", "fp8"] = "bf16", + revision: str | None = None, + local_files_only: bool = False, + enable_safety_checker: bool = True, +) -> StreamingCosmos3Pipeline: + """Build one auto-routing Cosmos pipeline for vision, sound, action, and transfer generation.""" + + if checkpoint_precision not in ("bf16", "fp8"): + raise ValueError(f"unsupported checkpoint precision: {checkpoint_precision}") + if checkpoint_precision == "fp8" and dtype != torch.bfloat16: + raise ValueError("the Cosmos 3 FP8 checkpoint requires BF16 compute inputs and outputs") + + if model_id is None: + model_id = DEFAULT_FP8_MODEL_ID if checkpoint_precision == "fp8" else "nvidia/Cosmos3-Nano" + component_model_id = model_id + component_revision = revision + fp8_transformer: torch.nn.Module | None = None + if checkpoint_precision == "fp8": + checkpoint_dir = materialize_fp8_checkpoint( + model_id, + revision=revision, + local_files_only=local_files_only, + ) + component_model_id = str(checkpoint_dir) + component_revision = None + fp8_transformer = load_fp8_transformer(checkpoint_dir, dtype=dtype) + + config = config or StreamingConfig() + runtime = StreamingCosmosRuntime(config) + blocks = Cosmos3OmniBlocks() + + denoise_loops = ( + _vision_denoise_loop(blocks), + _vision_sound_denoise_loop(blocks), + _vision_action_denoise_loop(blocks), + _vision_sound_action_denoise_loop(blocks), + ) + for denoise_loop in denoise_loops: + if "denoiser" not in denoise_loop.sub_blocks: + raise RuntimeError("a Cosmos denoise loop no longer exposes a 'denoiser' sub-block") + denoise_loop.sub_blocks["denoiser"] = StreamingCosmos3LoopDenoiser(runtime) + transfer_denoise_loop = _transfer_denoise_loop(blocks) + if "denoiser" not in transfer_denoise_loop.sub_blocks: + raise RuntimeError("the Cosmos transfer denoise loop no longer exposes a 'denoiser' sub-block") + transfer_denoise_loop.sub_blocks["denoiser"] = StreamingCosmos3TransferLoopDenoiser(runtime) + + if runtime.memory_trace is not None: + for denoise_loop in denoise_loops: + for nested_name in ("prepare_vision", "prepare_sound", "update_vision", "update_sound"): + if nested_name not in denoise_loop.sub_blocks: + continue + nested = denoise_loop.sub_blocks[nested_name] + denoise_loop.sub_blocks[nested_name] = _MemoryTracedBlock( + nested, + runtime, + f"denoise.{nested_name}", + ) + for nested_name in ("prepare_transfer", "update_transfer"): + nested = transfer_denoise_loop.sub_blocks[nested_name] + transfer_denoise_loop.sub_blocks[nested_name] = _MemoryTracedBlock( + nested, + runtime, + f"denoise.{nested_name}", + ) + for block_name in list(blocks.sub_blocks): + block = blocks.sub_blocks[block_name] + blocks.sub_blocks[block_name] = _MemoryTracedBlock(block, runtime, block_name) + + pipeline_kwargs: dict[str, Any] = {"local_files_only": local_files_only} + if component_revision is not None: + pipeline_kwargs["revision"] = component_revision + pipe = StreamingCosmos3Pipeline( + blocks=blocks, + pretrained_model_name_or_path=component_model_id, + streaming_runtime=runtime, + execution_device=config.device, + **pipeline_kwargs, + ) + component_load_kwargs: dict[str, Any] = { + "dtype": dtype, + "pretrained_model_name_or_path": component_model_id, + "local_files_only": local_files_only, + } + if component_revision is not None: + component_load_kwargs["revision"] = component_revision + if fp8_transformer is not None: + pipe.update_components(transformer=fp8_transformer) + component_paths, component_subfolders = _fp8_component_load_overrides(pipe, component_model_id) + component_load_kwargs["pretrained_model_name_or_path"] = component_paths + component_load_kwargs["subfolder"] = component_subfolders + # This superset loads the sound tokenizer once so calls can opt in with + # ``enable_sound=True`` without rebuilding a pipeline or changing graphs. + pipe.load_components(workflow="text2video_with_sound", **component_load_kwargs) + + original_vae_encode = pipe.vae.encode + original_vae_decode = pipe.vae.decode + + def managed_vae_encode(*args: Any, **kwargs: Any) -> Any: + if runtime.memory_trace is not None: + runtime.trace_checkpoint("vae_encode.start") + try: + return original_vae_encode(*args, **kwargs) + finally: + if runtime.memory_trace is not None: + runtime.trace_checkpoint("vae_encode.end") + + def managed_vae_decode(*args: Any, **kwargs: Any) -> Any: + runtime.prepare_for_decode() + if runtime.memory_trace is not None: + runtime.trace_checkpoint("vae_decode.start") + try: + return original_vae_decode(*args, **kwargs) + finally: + if runtime.memory_trace is not None: + runtime.trace_checkpoint("vae_decode.end") + + pipe.vae.encode = managed_vae_encode + pipe.vae.decode = managed_vae_decode + + original_decode_sound = pipe.decode_sound + + def managed_decode_sound(*args: Any, **kwargs: Any) -> Any: + """Run the official audio decoder on GPU, then return its weights to RAM.""" + + if pipe.sound_tokenizer is None: + raise RuntimeError("sound generation was requested but the checkpoint has no sound tokenizer") + runtime.prepare_for_decode() + runtime.trace_checkpoint("sound_decode.start") + pipe.sound_tokenizer.to(torch.device(config.device)) + try: + return original_decode_sound(*args, **kwargs) + finally: + pipe.sound_tokenizer.to(torch.device(config.offload_device)) + runtime.trace_checkpoint("sound_decode.end") + + pipe.decode_sound = managed_decode_sound + if enable_safety_checker: + pipe.enable_safety_checker() + else: + pipe.disable_safety_checker() + + if config.vae_leaf_offload: + pipe.vae.enable_group_offload( + onload_device=torch.device(config.device), + offload_device=torch.device(config.offload_device), + offload_type="leaf_level", + use_stream=False, + ) + device = torch.device(config.device) + total_vram = torch.cuda.get_device_properties(device).total_memory if device.type == "cuda" else 0 + vae_policy = _resolve_vae_memory_policy(config.memory_profile, total_vram) + runtime.resolve_memory_policy(total_vram) + runtime.vae_decode_mode = "tiled" if vae_policy.tile_size is not None else "untiled" + runtime.release_cache_before_safety = vae_policy.release_cache_before_safety + if vae_policy.tile_size is not None and vae_policy.tile_stride is not None: + pipe.vae.enable_tiling( + tile_sample_min_height=vae_policy.tile_size, + tile_sample_min_width=vae_policy.tile_size, + tile_sample_stride_height=vae_policy.tile_stride, + tile_sample_stride_width=vae_policy.tile_stride, + ) + return pipe + + +def build_text_to_video_pipeline(*args: Any, **kwargs: Any) -> StreamingCosmos3Pipeline: + """Compatibility alias for :func:`build_pipeline`. + + The returned pipeline is auto-routing and can also run T2I, I2V, V2V, + sound, action, and transfer workflows. + """ + + return build_pipeline(*args, **kwargs) diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py new file mode 100644 index 000000000..46c3adcea --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py @@ -0,0 +1,51 @@ +"""Lifecycle operations for the Diffusers group-offload hooks we install.""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch + + +def offload_groups(module: torch.nn.Module) -> None: + """Restore every group's CPU weights, including interrupted/prefetched groups.""" + from diffusers.hooks.group_offloading import GroupOffloadingHook + + groups = {} + for child in module.modules(): + registry = getattr(child, "_diffusers_hook", None) + if registry is not None: + for hook in registry.hooks.values(): + if isinstance(hook, GroupOffloadingHook): + groups[id(hook.group)] = hook.group + for group in groups.values(): + if group.stream is not None: + group.stream.synchronize() + device = torch.device(group.onload_device) + context = torch.cuda.device(device) if device.type == "cuda" else nullcontext() + with context: + group.offload_() + # The pinned streamed hook restores module parameters, but not module + # buffers. Preserve tensor identities used by its CPU backing store. + for child in group.modules: + for buffer in child.buffers(): + if buffer.device != group.offload_device: + backing = group.cpu_param_dict.get(buffer) + buffer.data = backing if backing is not None else buffer.data.to(group.offload_device) + + +def remove_group_hooks(module: torch.nn.Module) -> None: + """Remove only offload/prefetch hooks; preserve other hooks and FP8 wrappers.""" + from diffusers.hooks.group_offloading import ( + GroupOffloadingHook, + LayerExecutionTrackerHook, + LazyPrefetchGroupOffloadingHook, + ) + + hook_types = (GroupOffloadingHook, LayerExecutionTrackerHook, LazyPrefetchGroupOffloadingHook) + for child in module.modules(): + registry = getattr(child, "_diffusers_hook", None) + if registry is not None: + for name, hook in reversed(list(registry.hooks.items())): + if isinstance(hook, hook_types): + registry.remove_hook(name, recurse=False) diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py new file mode 100644 index 000000000..4af34f21a --- /dev/null +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py @@ -0,0 +1,970 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from .attention.loader import load_attention_backend +from .block_views import ( + MEMORY_SERIAL_MLP_CHUNK_TOKENS, + GenerationStack, + PrefixKV, + RotaryPair, + UnderstandingStack, +) +from .config import StreamingConfig +from .memory_trace import CudaMemoryTrace +from .metrics import RuntimeStats +from .offloading import offload_groups, remove_group_hooks + +_GIB = 2**30 +_AUTO_PREFETCH_MIN_VRAM_BYTES = 11 * _GIB +_AUTO_PREFETCH_UNCONDITIONAL_VRAM_BYTES = 14 * _GIB +_AUTO_PREFETCH_LARGE_BUFFER_MIN_VRAM_BYTES = 20 * _GIB +# Three 480p Cosmos3-Nano transfer branches occupy about 490 MiB in their +# persistent generation hidden states. The 640 MiB ceiling leaves the 12 GiB +# card enough room for the additional decoder block held by prefetch, while +# keeping 720p on the lower-peak synchronous path. +_AUTO_PREFETCH_MAX_GENERATION_STATE_BYTES = 640 * 2**20 +_AUTO_BRANCH_STAGING_MAX_VRAM_BYTES = 14 * _GIB +_AUTO_BRANCH_STAGING_MIN_STATE_BYTES = _AUTO_PREFETCH_MAX_GENERATION_STATE_BYTES +_TRANSFER_BRANCH_NAMES = {"cond_full", "cond_no_control", "uncond_full"} + + +@dataclass(slots=True) +class _BranchCache: + und_len: int + generation_rotary: RotaryPair + prefixes: list[PrefixKV] + + +@dataclass(slots=True) +class _PreparedGeneration: + hidden: torch.Tensor + original_latent_shapes: list[tuple[int, int, int]] + + +@dataclass(slots=True) +class _ModalityPredictions: + vision: list[torch.Tensor] + sound: list[torch.Tensor] | None + action: list[torch.Tensor] | None + + +@dataclass(frozen=True, slots=True) +class _GenerationPrefetchPolicy: + enabled: bool + reason: str + branch_count: int + generation_tokens: int + estimated_state_bytes: int + + +@dataclass(frozen=True, slots=True) +class _GenerationBranchStagingPolicy: + enabled: bool + reason: str + branch_count: int + estimated_state_bytes: int + + +class StreamingCosmosRuntime: + """Request-scoped Cosmos 3 understanding cache and split executor. + + Supports the coupled vision, sound, and action denoising paths. Transfer + chunk scheduling is integrated by the modular-pipeline adapter. + """ + + def __init__(self, config: StreamingConfig) -> None: + self.config = config + self.device = torch.device(config.device) + self.offload_device = torch.device(config.offload_device) + self.backend = load_attention_backend(config.attention_backend, **config.attention_backend_kwargs) + self.stats = RuntimeStats() + self.release_cache_before_safety = False + self.detected_vram_bytes = 0 + self.effective_use_stream = config.use_stream + self.generation_prefetch_reason = "unresolved" + self.generation_workload_tokens = 0 + self.estimated_generation_state_bytes = 0 + self.generation_branch_staging = False + self.generation_branch_staging_reason = "unresolved" + self.vae_decode_mode = "unresolved" + self.memory_trace = ( + CudaMemoryTrace( + config.memory_trace_path, + self.device, + synchronize=config.memory_trace_synchronize, + ) + if config.memory_trace_path is not None + else None + ) + + self._active = False + self._transformer: torch.nn.Module | None = None + self._understanding_stack: UnderstandingStack | None = None + self._generation_stack: GenerationStack | None = None + self._branch_cache: dict[str, _BranchCache] = {} + self._cache_signature: tuple[Any, ...] | None = None + self._understanding_signature: tuple[Any, ...] | None = None + self._precision_signature: tuple[bool, bool] | None = None + self._inference_tensor_versions: dict[int, tuple[torch.Tensor, torch.Tensor, int]] = {} + self._generation_modules_resident = False + self._decode_prepared = False + self._installed_generation_use_stream: bool | None = None + self._request_prefetch_signature: tuple[Any, ...] | None = None + self._request_branch_staging_signature: tuple[Any, ...] | None = None + + def begin_request(self) -> None: + if self._active: + raise RuntimeError("StreamingCosmosRuntime is not reentrant") + self._active = True + self._precision_signature = None + self._decode_prepared = False + self._branch_cache.clear() + self._inference_tensor_versions.clear() + self._cache_signature = None + self._understanding_signature = None + self._request_prefetch_signature = None + self._request_branch_staging_signature = None + self.stats = RuntimeStats( + understanding_weight_bytes=self.stats.understanding_weight_bytes, + generation_weight_bytes=self.stats.generation_weight_bytes, + ) + if self.memory_trace is not None: + try: + self.memory_trace.begin() + except BaseException: + self._active = False + raise + self.trace_checkpoint( + "memory_policy", + configured_profile=self.config.memory_profile, + detected_vram_bytes=self.detected_vram_bytes, + vae_decode_mode=self.vae_decode_mode, + generation_use_stream=self.effective_use_stream, + ) + + def resolve_memory_policy(self, total_vram_bytes: int) -> None: + """Resolve backend-specific internals from the intent-level profile.""" + + self.detected_vram_bytes = total_vram_bytes + uses_large_buffers = bool(getattr(self.backend, "retains_large_device_buffers", False)) + sub_14gb_auto = self.config.memory_profile == "auto" and total_vram_bytes < 14 * 2**30 + constrained_auto = ( + self.config.memory_profile == "auto" and total_vram_bytes < _AUTO_PREFETCH_LARGE_BUFFER_MIN_VRAM_BYTES + ) + self.effective_use_stream = self.config.use_stream and not ( + sub_14gb_auto or (uses_large_buffers and constrained_auto) + ) + + def end_request(self, *, failed: bool = False) -> None: + try: + self._branch_cache.clear() + self._inference_tensor_versions.clear() + self._cache_signature = None + self._understanding_signature = None + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(self.device) + for stack in (self._understanding_stack, self._generation_stack): + if stack is not None and self.config.enable_group_offload: + offload_groups(stack) + if failed: + remove_group_hooks(stack) + self._offload_generation_modules() + if failed and self._transformer is not None: + self._transformer.to(self.offload_device) + # A failed first forward may leave a partial lazy-prefetch + # trace. Rebuild views/hooks on retry instead of reusing it. + self._transformer = None + self._understanding_stack = None + self._generation_stack = None + self._installed_generation_use_stream = None + release_buffers = getattr(self.backend, "release_buffers", None) + if release_buffers is not None and not self._decode_prepared: + release_buffers() + if self.device.type == "cuda" and torch.cuda.is_available() and self.config.empty_cuda_cache_on_end: + torch.cuda.empty_cache() + finally: + try: + if self.memory_trace is not None: + self.memory_trace.end() + finally: + self._active = False + + def prepare_for_decode(self) -> None: + """Release denoising-only state before VAE decode begins.""" + + if self._decode_prepared: + return + self.trace_checkpoint("denoising_resources.release.start") + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(self.device) + self._branch_cache.clear() + self._inference_tensor_versions.clear() + self._cache_signature = None + self._understanding_signature = None + self._offload_generation_modules() + release_buffers = getattr(self.backend, "release_buffers", None) + if release_buffers is not None: + release_buffers() + if self.device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.empty_cache() + self._decode_prepared = True + self.trace_checkpoint("denoising_resources.release.end") + + @torch.no_grad() + def predict_modalities( + self, + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> dict[str, _ModalityPredictions]: + """Predict per-branch vision, sound, and action velocities.""" + + if not self._active: + raise RuntimeError("begin_request() must be called before predict_modalities()") + if not branches: + raise ValueError("at least one prediction branch is required") + # Transfer decodes every autoregressive chunk. A new denoising phase + # after that decode must be allowed to release its own caches again. + if self._decode_prepared: + self._decode_prepared = False + + self._bind_transformer(transformer, branches) + self._invalidate_changed_precision(transformer) + self.stats.denoising_calls += 1 + denoising_call = self.stats.denoising_calls - 1 + self.trace_checkpoint("transformer.start", step=denoising_call) + self.stats.estimated_block_h2d_bytes += self.stats.generation_weight_bytes + signature = self._static_signature(branches) + if not self._branch_cache: + self._prefill_understanding(branches) + self._cache_signature = signature + elif signature != self._cache_signature: + understanding_signature = self._static_understanding_signature(branches) + if understanding_signature == self._understanding_signature and set(branches) == set(self._branch_cache): + self._refresh_generation_rotaries(branches) + else: + self._branch_cache.clear() + self._prefill_understanding(branches) + self._cache_signature = signature + + self._onload_generation_modules() + branch_policy = self._resolve_generation_branch_staging_policy(transformer, branches) + self._record_generation_branch_staging_policy(branch_policy) + prepared: dict[str, _PreparedGeneration] = {} + for name, kwargs in branches.items(): + item = self._prepare_generation(kwargs) + if branch_policy.enabled: + item.hidden = self._stage_hidden_on_host(item.hidden) + prepared[name] = item + if branch_policy.enabled and self.device.type == "cuda": + # Preparation uses differently sized temporary tensors. Return those + # cached segments before the first large MLP allocation so the + # serial branch working set starts from a compact allocator state. + torch.cuda.synchronize(self.device) + torch.cuda.empty_cache() + branch_names = list(branches) + first_name = branch_names[0] + second_name = branch_names[1] if len(branch_names) > 1 else None + extra_names = branch_names[2:] + first = prepared[first_name] + second = None if second_name is None else prepared[second_name] + first_cache = self._branch_cache[first_name] + second_cache = None if second_name is None else self._branch_cache[second_name] + assert self._generation_stack is not None + if extra_names: + first_hidden, second_hidden, extra_hiddens = self._generation_stack( + first.hidden, + first_cache.generation_rotary, + first_cache.prefixes, + None if second is None else second.hidden, + None if second_cache is None else second_cache.generation_rotary, + None if second_cache is None else second_cache.prefixes, + [prepared[name].hidden for name in extra_names], + [self._branch_cache[name].generation_rotary for name in extra_names], + [self._branch_cache[name].prefixes for name in extra_names], + stage_branches=branch_policy.enabled, + execution_device=self.device, + offload_device=self.offload_device, + ) + else: + first_hidden, second_hidden = self._generation_stack( + first.hidden, + first_cache.generation_rotary, + first_cache.prefixes, + None if second is None else second.hidden, + None if second_cache is None else second_cache.generation_rotary, + None if second_cache is None else second_cache.prefixes, + stage_branches=branch_policy.enabled, + execution_device=self.device, + offload_device=self.offload_device, + ) + extra_hiddens = [] + + outputs = { + first_name: self._decode_predictions(first_hidden, branches[first_name], first.original_latent_shapes) + } + if second_name is not None: + if second_hidden is None or second is None: + raise RuntimeError("generation stack did not return its second branch") + outputs[second_name] = self._decode_predictions( + second_hidden, + branches[second_name], + second.original_latent_shapes, + ) + for name, hidden in zip(extra_names, extra_hiddens): + outputs[name] = self._decode_predictions( + hidden, + branches[name], + prepared[name].original_latent_shapes, + ) + self.trace_checkpoint("transformer.end", step=denoising_call) + return outputs + + def predict_vision( + self, + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> dict[str, list[torch.Tensor]]: + """Compatibility wrapper for callers that only consume vision.""" + + return {name: prediction.vision for name, prediction in self.predict_modalities(transformer, branches).items()} + + def _invalidate_changed_precision(self, transformer: torch.nn.Module) -> None: + from diffusers.pipelines.cosmos import mixed_precision as mixed + + native = getattr(transformer, mixed._RUNTIME_ATTRIBUTE, None) + signature = ( + (native.use_high_precision("reasoner"), native.use_high_precision("generation")) + if native is not None and native.active + else None + ) + if signature != self._precision_signature: + # Native dispatch may change understanding projections at a step + # boundary. K/V computed with the previous precision cannot be reused. + self._branch_cache.clear() + self._cache_signature = None + self._understanding_signature = None + self._precision_signature = signature + + def _bind_transformer( + self, + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> None: + policy = self._resolve_generation_prefetch_policy(transformer, branches) + if self._transformer is transformer: + self._use_installed_prefetch_policy(policy) + return + if self._transformer is not None: + raise RuntimeError("a StreamingCosmosRuntime cannot be rebound to another transformer") + + self._transformer = transformer + block_callback = self._trace_block if self.config.memory_trace_blocks else None + self._understanding_stack = UnderstandingStack(transformer.layers, self.backend, block_callback) + self._generation_stack = GenerationStack(transformer.layers, self.backend, block_callback) + self.stats.understanding_weight_bytes = self._parameter_bytes(self._understanding_stack) + self.stats.generation_weight_bytes = self._parameter_bytes(self._generation_stack) + self.effective_use_stream = policy.enabled + self._installed_generation_use_stream = policy.enabled if self.config.enable_group_offload else False + self._record_generation_prefetch_policy(policy, actual_enabled=self.effective_use_stream) + + if self.config.enable_group_offload: + from diffusers.hooks import apply_group_offloading + + common = { + "onload_device": self.device, + "offload_device": self.offload_device, + "offload_type": "block_level", + "num_blocks_per_group": 1, + "low_cpu_mem_usage": self.config.low_cpu_mem_usage, + } + # Understanding is invoked once, so Diffusers' lazy streamed trace + # cannot benefit it. Synchronous block offload avoids a useless trace. + apply_group_offloading( + self._understanding_stack, + use_stream=False, + record_stream=False, + **common, + ) + self._install_generation_offload(self.effective_use_stream) + else: + self._understanding_stack.to(self.device) + self._generation_stack.to(self.device) + + def _resolve_generation_prefetch_policy( + self, + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> _GenerationPrefetchPolicy: + """Select async generation-block prefetch from the actual request shape. + + VAE tiling has a different lifetime and is resolved independently by + the modular-pipeline adapter. This policy only budgets the extra model + block that Diffusers keeps on the GPU while the current block computes. + """ + + branch_count = self._prefetch_branch_count(branches) + generation_tokens = self._generation_token_count(branches, branch_count) + # The official FP8 checkpoint keeps the persistent hidden state in + # BF16 and quantizes Linear inputs internally. Using FP8 weight size + # here would under-budget the activation state and could select an + # unsafe prefetch policy on a 12 GiB card. + dtype_bytes = self._generation_activation_element_size(transformer, branches) + hidden_size = int(transformer.config.hidden_size) + estimated_state_bytes = generation_tokens * hidden_size * dtype_bytes + + common = { + "branch_count": branch_count, + "generation_tokens": generation_tokens, + "estimated_state_bytes": estimated_state_bytes, + } + if not self.config.enable_group_offload: + return _GenerationPrefetchPolicy(False, "group_offload_disabled", **common) + if not self.config.use_stream: + return _GenerationPrefetchPolicy(False, "streaming_disabled", **common) + if self.config.memory_profile == "low_memory": + return _GenerationPrefetchPolicy(False, "low_memory_profile", **common) + if self.config.memory_profile == "performance": + return _GenerationPrefetchPolicy(True, "performance_profile", **common) + + uses_large_buffers = bool(getattr(self.backend, "retains_large_device_buffers", False)) + if uses_large_buffers and self.detected_vram_bytes < _AUTO_PREFETCH_LARGE_BUFFER_MIN_VRAM_BYTES: + return _GenerationPrefetchPolicy(False, "attention_backend_reserve", **common) + if self.detected_vram_bytes == 0: + return _GenerationPrefetchPolicy(self.config.use_stream, "vram_unknown", **common) + if self.detected_vram_bytes >= _AUTO_PREFETCH_UNCONDITIONAL_VRAM_BYTES: + return _GenerationPrefetchPolicy(True, "vram_headroom", **common) + if self.detected_vram_bytes < _AUTO_PREFETCH_MIN_VRAM_BYTES: + return _GenerationPrefetchPolicy(False, "insufficient_vram", **common) + if estimated_state_bytes <= _AUTO_PREFETCH_MAX_GENERATION_STATE_BYTES: + return _GenerationPrefetchPolicy(True, "workload_headroom", **common) + return _GenerationPrefetchPolicy(False, "workload_too_large", **common) + + def _resolve_generation_branch_staging_policy( + self, + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> _GenerationBranchStagingPolicy: + """Select host staging for Transfer branch hidden states. + + Transfer can evaluate as many as three guidance branches. The decoder + block already computes them serially while its weights are resident, + but the normal path retains every branch hidden state on CUDA. On a + 12 GiB card at 720p those persistent BF16 states consume enough memory + to prevent the generation MLP from allocating its temporary output. + """ + + branch_count = self._prefetch_branch_count(branches) + generation_tokens = self._generation_token_count(branches, branch_count) + dtype_bytes = self._generation_activation_element_size(transformer, branches) + estimated_state_bytes = generation_tokens * int(transformer.config.hidden_size) * dtype_bytes + common = {"branch_count": branch_count, "estimated_state_bytes": estimated_state_bytes} + + is_transfer = bool(_TRANSFER_BRANCH_NAMES.intersection(branches)) + if not is_transfer: + return _GenerationBranchStagingPolicy(False, "not_transfer", **common) + if len(branches) <= 1: + return _GenerationBranchStagingPolicy(False, "single_active_branch", **common) + if self.config.memory_profile == "performance": + return _GenerationBranchStagingPolicy(False, "performance_profile", **common) + if self.config.memory_profile == "low_memory": + return _GenerationBranchStagingPolicy(True, "low_memory_profile", **common) + if self.detected_vram_bytes == 0: + return _GenerationBranchStagingPolicy(False, "vram_unknown", **common) + if self.detected_vram_bytes >= _AUTO_BRANCH_STAGING_MAX_VRAM_BYTES: + return _GenerationBranchStagingPolicy(False, "vram_headroom", **common) + if estimated_state_bytes <= _AUTO_BRANCH_STAGING_MIN_STATE_BYTES: + return _GenerationBranchStagingPolicy(False, "workload_headroom", **common) + return _GenerationBranchStagingPolicy(True, "workload_pressure", **common) + + def _record_generation_branch_staging_policy(self, policy: _GenerationBranchStagingPolicy) -> None: + signature = (policy.enabled, policy.reason, policy.branch_count, policy.estimated_state_bytes) + if signature == self._request_branch_staging_signature: + return + self._request_branch_staging_signature = signature + self.generation_branch_staging = policy.enabled + self.generation_branch_staging_reason = policy.reason + self.trace_checkpoint( + "generation_branch_staging_policy", + enabled=policy.enabled, + reason=policy.reason, + branch_count=policy.branch_count, + estimated_generation_state_bytes=policy.estimated_state_bytes, + mlp_chunk_tokens=MEMORY_SERIAL_MLP_CHUNK_TOKENS if policy.enabled else None, + detected_vram_bytes=self.detected_vram_bytes, + ) + + @staticmethod + def _prefetch_branch_count(branches: dict[str, dict[str, Any]]) -> int: + # Transfer guidance intervals can change which branches execute during + # a request. Budget all three branches up front so a later interval does + # not invalidate hooks installed at the first denoising step. + if _TRANSFER_BRANCH_NAMES.intersection(branches): + return 3 + return len(branches) + + @staticmethod + def _generation_token_count(branches: dict[str, dict[str, Any]], branch_count: int) -> int: + lengths = [max(0, int(kwargs["sequence_length"]) - int(kwargs["und_len"])) for kwargs in branches.values()] + if not lengths: + return 0 + if branch_count == len(lengths): + return sum(lengths) + return max(lengths) * branch_count + + @staticmethod + def _generation_activation_element_size( + transformer: torch.nn.Module, + branches: dict[str, dict[str, Any]], + ) -> int: + for kwargs in branches.values(): + for field in ("vision_tokens", "sound_tokens", "action_tokens"): + tensors = kwargs.get(field) + if tensors: + return tensors[0].element_size() + config_dtype = getattr(transformer.config, "dtype", "bfloat16") + if isinstance(config_dtype, torch.dtype): + dtype = config_dtype + else: + dtype = getattr(torch, str(config_dtype).removeprefix("torch."), torch.bfloat16) + return torch.tensor([], dtype=dtype).element_size() + + def _use_installed_prefetch_policy(self, policy: _GenerationPrefetchPolicy) -> None: + installed = self._installed_generation_use_stream + if installed is None: + raise RuntimeError("generation offload hooks were not initialized") + if installed != policy.enabled and self.config.enable_group_offload: + # No generation forward is active here. Drain the previous stream + # before replacing its hooks and pinned CPU backing stores. A new + # async policy gets a fresh lazy execution trace on its first call. + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + stack = self._require_generation_stack() + offload_groups(stack) + remove_group_hooks(stack) + self._install_generation_offload(policy.enabled) + self._installed_generation_use_stream = policy.enabled + self.effective_use_stream = policy.enabled + self._record_generation_prefetch_policy(policy, actual_enabled=policy.enabled) + + def _install_generation_offload(self, use_stream: bool) -> None: + from diffusers.hooks import apply_group_offloading + + apply_group_offloading( + self._require_generation_stack(), + onload_device=self.device, + offload_device=self.offload_device, + offload_type="block_level", + num_blocks_per_group=1, + low_cpu_mem_usage=self.config.low_cpu_mem_usage, + use_stream=use_stream, + record_stream=self.config.record_stream and use_stream, + ) + + def _require_generation_stack(self) -> GenerationStack: + if self._generation_stack is None: + raise RuntimeError("generation stack is not initialized") + return self._generation_stack + + def _record_generation_prefetch_policy( + self, + policy: _GenerationPrefetchPolicy, + *, + actual_enabled: bool, + ) -> None: + signature = ( + actual_enabled, + policy.reason, + policy.branch_count, + policy.generation_tokens, + policy.estimated_state_bytes, + ) + if signature == self._request_prefetch_signature: + return + self._request_prefetch_signature = signature + self.generation_prefetch_reason = policy.reason + self.generation_workload_tokens = policy.generation_tokens + self.estimated_generation_state_bytes = policy.estimated_state_bytes + self.trace_checkpoint( + "generation_prefetch_policy", + generation_use_stream=actual_enabled, + reason=policy.reason, + branch_count=policy.branch_count, + generation_tokens=policy.generation_tokens, + estimated_generation_state_bytes=policy.estimated_state_bytes, + max_auto_generation_state_bytes=_AUTO_PREFETCH_MAX_GENERATION_STATE_BYTES, + detected_vram_bytes=self.detected_vram_bytes, + ) + + def _tensor_signature(self, tensor: torch.Tensor) -> tuple[Any, ...]: + if not torch.is_inference(tensor): + return (id(tensor), tensor._version, tuple(tensor.shape)) + # Inference tensors have no version counter. Keep a request-local + # snapshot so in-place edits still invalidate cached conditioning. + # Only this fallback needs a content comparison (and potentially a sync). + previous = self._inference_tensor_versions.get(id(tensor)) + revision = 0 if previous is None else previous[2] + if previous is None or not torch.equal(tensor, previous[1]): + revision += 1 + with torch.inference_mode(False): + snapshot = tensor.detach().clone() + self._inference_tensor_versions[id(tensor)] = (tensor, snapshot, revision) + return (id(tensor), revision, tuple(tensor.shape)) + + def _static_signature(self, branches: dict[str, dict[str, Any]]) -> tuple[Any, ...]: + """Cheap identity signature for modular pipeline fields that must stay static. + + The official Cosmos modular workflow reuses these tensor objects through + the denoising loop. Avoiding a content hash prevents a GPU-to-CPU sync on + every diffusion step. + """ + + signature: list[Any] = [] + for name, kwargs in branches.items(): + signature.extend((name, int(kwargs["und_len"]), int(kwargs["sequence_length"]))) + for field in ("input_ids", "text_indexes", "position_ids"): + tensor = kwargs[field] + signature.append((field, self._tensor_signature(tensor))) + return tuple(signature) + + def _static_understanding_signature(self, branches: dict[str, dict[str, Any]]) -> tuple[Any, ...]: + signature: list[Any] = [] + for name, kwargs in branches.items(): + signature.extend((name, int(kwargs["und_len"]))) + for field in ("input_ids", "text_indexes"): + tensor = kwargs[field] + signature.append((field, self._tensor_signature(tensor))) + signature.append(self._understanding_positions_key(kwargs)) + return tuple(signature) + + @staticmethod + def _understanding_positions_key(kwargs: dict[str, Any]) -> tuple[Any, ...]: + # Compare prefix contents, not the full position tensor's version: + # Transfer may change generation positions without changing its prompt. + # Called only during prefill or when the cheap static signature changes, + # never on the ordinary unchanged-tensor denoising path. + prefix = kwargs["position_ids"][..., : int(kwargs["und_len"])] + return (tuple(prefix.shape), tuple(prefix.reshape(-1).tolist())) + + def _prefill_understanding(self, branches: dict[str, dict[str, Any]]) -> None: + transformer = self._require_transformer() + understanding_stack = self._require_understanding_stack() + self.stats.understanding_prefills += 1 + understanding_groups: dict[tuple[Any, ...], list[str]] = {} + for name, kwargs in branches.items(): + key = self._understanding_branch_key(kwargs) + understanding_groups.setdefault(key, []).append(name) + self.stats.estimated_block_h2d_bytes += self.stats.understanding_weight_bytes * ( + (len(understanding_groups) + 1) // 2 + ) + self.trace_checkpoint("understanding_prefill.start") + + embed_tokens = transformer.embed_tokens + embed_tokens.to(self.device) + try: + group_items = list(understanding_groups.items()) + prefixes_by_key: dict[tuple[Any, ...], list[PrefixKV]] = {} + for group_index in range(0, len(group_items), 2): + first_key, first_names = group_items[group_index] + first_kwargs = branches[first_names[0]] + first_hidden = self._prepare_understanding_hidden(first_kwargs) + first_rotary, _ = self._compute_rotary(first_kwargs) + + second_key = None + second_hidden = None + second_rotary = None + if group_index + 1 < len(group_items): + second_key, second_names = group_items[group_index + 1] + second_kwargs = branches[second_names[0]] + second_hidden = self._prepare_understanding_hidden(second_kwargs) + second_rotary, _ = self._compute_rotary(second_kwargs) + + first_prefixes, second_prefixes = understanding_stack( + first_hidden, + first_rotary, + second_hidden, + second_rotary, + ) + prefixes_by_key[first_key] = first_prefixes + if second_key is not None: + if second_prefixes is None: + raise RuntimeError("understanding stack did not produce its second branch cache") + prefixes_by_key[second_key] = second_prefixes + finally: + embed_tokens.to(self.offload_device) + + for name, kwargs in branches.items(): + _, generation_rotary = self._compute_rotary(kwargs) + self._branch_cache[name] = _BranchCache( + und_len=int(kwargs["und_len"]), + generation_rotary=generation_rotary, + prefixes=prefixes_by_key[self._understanding_branch_key(kwargs)], + ) + self._understanding_signature = self._static_understanding_signature(branches) + self.trace_checkpoint("understanding_prefill.end") + + def _understanding_branch_key(self, kwargs: dict[str, Any]) -> tuple[Any, ...]: + return ( + int(kwargs["und_len"]), + self._tensor_signature(kwargs["input_ids"]), + self._tensor_signature(kwargs["text_indexes"]), + self._understanding_positions_key(kwargs), + ) + + def _refresh_generation_rotaries(self, branches: dict[str, dict[str, Any]]) -> None: + for name, kwargs in branches.items(): + _, generation_rotary = self._compute_rotary(kwargs) + cache = self._branch_cache[name] + cache.generation_rotary = generation_rotary + + def trace_checkpoint(self, event: str, **fields: Any) -> None: + if self.memory_trace is not None: + self.memory_trace.checkpoint(event, **fields) + + def _trace_block(self, tower: str, block: int) -> None: + self.trace_checkpoint( + f"{tower}.block.end", + step=max(self.stats.denoising_calls - 1, 0), + block=block, + ) + + def _prepare_understanding_hidden(self, kwargs: dict[str, Any]) -> torch.Tensor: + transformer = self._require_transformer() + input_ids = kwargs["input_ids"].to(self.device) + text_indexes = kwargs["text_indexes"].to(self.device) + und_len = int(kwargs["und_len"]) + + embeddings = transformer.embed_tokens(input_ids) + hidden = embeddings.new_zeros((und_len, transformer.config.hidden_size)) + if text_indexes.numel() and (text_indexes.min() < 0 or text_indexes.max() >= und_len): + raise ValueError("text indexes must address only the understanding prefix") + hidden[text_indexes] = embeddings + return hidden + + def _compute_rotary(self, kwargs: dict[str, Any]) -> tuple[RotaryPair, RotaryPair]: + transformer = self._require_transformer() + position_ids = kwargs["position_ids"].to(self.device) + und_len = int(kwargs["und_len"]) + + # Keep this shape handling identical to Cosmos3OmniTransformer.forward. + rotary_position_ids = position_ids.unsqueeze(0) if position_ids.ndim == 1 else position_ids.unsqueeze(1) + parameter = next(transformer.embed_tokens.parameters()) + cos, sin = transformer.rotary_emb( + position_ids=rotary_position_ids, + device=self.device, + dtype=parameter.dtype, + ) + cos = cos.squeeze(0) + sin = sin.squeeze(0) + return (cos[:und_len], sin[:und_len]), (cos[und_len:], sin[und_len:]) + + def _prepare_generation(self, kwargs: dict[str, Any]) -> _PreparedGeneration: + transformer = self._require_transformer() + vision_tokens = [tensor.to(device=self.device, dtype=transformer.dtype) for tensor in kwargs["vision_tokens"]] + packed, original_shapes = transformer._patchify_and_pack_latents(vision_tokens) + packed = transformer.proj_in(packed) + + timesteps = kwargs["vision_timesteps"].to(self.device) * transformer.config.timestep_scale + time_dtype = next(transformer.time_embedder.parameters()).dtype + timestep_embeddings = transformer.time_embedder(transformer.time_proj(timesteps).to(time_dtype)) + timestep_embeddings = timestep_embeddings.to(packed.dtype) + packed = transformer._apply_timestep_embeds_to_noisy_tokens( + packed_tokens=packed, + packed_timestep_embeds=timestep_embeddings, + noisy_frame_indexes=kwargs["vision_noisy_frame_indexes"], + token_shapes=kwargs["vision_token_shapes"], + ) + + und_len = int(kwargs["und_len"]) + generation_length = int(kwargs["sequence_length"]) - und_len + hidden = packed.new_zeros((generation_length, transformer.config.hidden_size)) + local_indexes = kwargs["vision_sequence_indexes"].to(self.device) - und_len + if local_indexes.numel() and (local_indexes.min() < 0 or local_indexes.max() >= generation_length): + raise ValueError("vision indexes must address only the generation suffix") + hidden[local_indexes] = packed + + has_sound = "sound_tokens" in kwargs + if has_sound: + required_sound = { + "sound_token_shapes", + "sound_sequence_indexes", + "sound_timesteps", + "sound_noisy_frame_indexes", + } + missing_sound = sorted(required_sound - kwargs.keys()) + if missing_sound: + raise KeyError(f"missing Cosmos sound transformer inputs: {missing_sound}") + if not hasattr(transformer, "audio_proj_in"): + raise RuntimeError("the loaded Cosmos transformer does not include the sound projection heads") + + sound_tokens = [tensor.to(device=self.device, dtype=transformer.dtype) for tensor in kwargs["sound_tokens"]] + packed_sound = transformer._pack_sound_latents(sound_tokens, kwargs["sound_token_shapes"]) + packed_sound = transformer.audio_proj_in(packed_sound) + packed_sound = packed_sound + transformer.audio_modality_embed.to( + device=self.device, dtype=packed_sound.dtype + ) + sound_timesteps = kwargs["sound_timesteps"].to(self.device) * transformer.config.timestep_scale + sound_timestep_embeddings = transformer.time_embedder( + transformer.time_proj(sound_timesteps).to(time_dtype) + ).to(packed_sound.dtype) + packed_sound = transformer._apply_timestep_embeds_to_noisy_tokens( + packed_tokens=packed_sound, + packed_timestep_embeds=sound_timestep_embeddings, + noisy_frame_indexes=kwargs["sound_noisy_frame_indexes"], + token_shapes=kwargs["sound_token_shapes"], + ) + sound_indexes = kwargs["sound_sequence_indexes"].to(self.device) - und_len + if sound_indexes.numel() and (sound_indexes.min() < 0 or sound_indexes.max() >= generation_length): + raise ValueError("sound indexes must address only the generation suffix") + hidden[sound_indexes] = packed_sound + + has_action = "action_tokens" in kwargs + if has_action: + required_action = { + "action_token_shapes", + "action_sequence_indexes", + "action_mse_loss_indexes", + "action_timesteps", + "action_noisy_frame_indexes", + "action_domain_ids", + } + missing_action = sorted(required_action - kwargs.keys()) + if missing_action: + raise KeyError(f"missing Cosmos action transformer inputs: {missing_action}") + if not hasattr(transformer, "action_proj_in"): + raise RuntimeError("the loaded Cosmos transformer does not include the action projection heads") + + action_tokens = [ + tensor.to(device=self.device, dtype=transformer.dtype) for tensor in kwargs["action_tokens"] + ] + packed_action, action_domain_ids = transformer._pack_action_latents( + action_tokens, + kwargs["action_token_shapes"], + kwargs["action_domain_ids"], + ) + action_domain_ids = action_domain_ids.to(self.device) + packed_action = transformer.action_proj_in(packed_action, action_domain_ids) + packed_action = packed_action + transformer.action_modality_embed.to( + device=self.device, dtype=packed_action.dtype + ) + if kwargs["action_mse_loss_indexes"].numel() > 0: + action_timesteps = kwargs["action_timesteps"].to(self.device) * transformer.config.timestep_scale + action_timestep_embeddings = transformer.time_embedder( + transformer.time_proj(action_timesteps).to(time_dtype) + ).to(packed_action.dtype) + packed_action = transformer._apply_timestep_embeds_to_noisy_tokens( + packed_tokens=packed_action, + packed_timestep_embeds=action_timestep_embeddings, + noisy_frame_indexes=kwargs["action_noisy_frame_indexes"], + token_shapes=kwargs["action_token_shapes"], + ) + action_indexes = kwargs["action_sequence_indexes"].to(self.device) - und_len + if action_indexes.numel() and (action_indexes.min() < 0 or action_indexes.max() >= generation_length): + raise ValueError("action indexes must address only the generation suffix") + hidden[action_indexes] = packed_action + return _PreparedGeneration(hidden=hidden, original_latent_shapes=original_shapes) + + def _stage_hidden_on_host(self, hidden: torch.Tensor) -> torch.Tensor: + if hidden.device == self.offload_device: + return hidden + if self.offload_device.type == "cpu" and self.device.type == "cuda": + staged = torch.empty_like(hidden, device=self.offload_device, pin_memory=True) + staged.copy_(hidden, non_blocking=True) + return staged + return hidden.to(self.offload_device) + + @staticmethod + def _parameter_bytes(module: torch.nn.Module) -> int: + unique: dict[int, torch.nn.Parameter] = {id(parameter): parameter for parameter in module.parameters()} + return sum(parameter.numel() * parameter.element_size() for parameter in unique.values()) + + def _decode_predictions( + self, + hidden: torch.Tensor, + kwargs: dict[str, Any], + original_latent_shapes: list[tuple[int, int, int]], + ) -> _ModalityPredictions: + transformer = self._require_transformer() + if hidden.device != self.device: + hidden = hidden.to( + self.device, + non_blocking=hidden.device.type == "cpu" and hidden.is_pinned(), + ) + hidden = transformer.norm_moe_gen(hidden) + indexes = kwargs["vision_mse_loss_indexes"].to(self.device) - int(kwargs["und_len"]) + packed_predictions = transformer.proj_out(hidden[indexes]) + vision = transformer._unpatchify_and_unpack_latents( + packed_predictions, + token_shapes_vision=kwargs["vision_token_shapes"], + noisy_frame_indexes_vision=kwargs["vision_noisy_frame_indexes"], + original_latent_shapes=original_latent_shapes, + ) + sound = None + if "sound_tokens" in kwargs: + indexes = kwargs["sound_mse_loss_indexes"].to(self.device) - int(kwargs["und_len"]) + packed_sound_predictions = transformer.audio_proj_out(hidden[indexes]) + sound = transformer._unpack_sound_latents( + packed_sound_predictions, + kwargs["sound_token_shapes"], + kwargs["sound_noisy_frame_indexes"], + ) + action = None + if "action_tokens" in kwargs: + per_noisy_domain_ids = [ + domain_id.reshape(1).expand(len(noisy_indexes)) + for domain_id, noisy_indexes in zip(kwargs["action_domain_ids"], kwargs["action_noisy_frame_indexes"]) + ] + per_noisy_domain_ids = torch.cat(per_noisy_domain_ids).to(self.device) + indexes = kwargs["action_mse_loss_indexes"].to(self.device) - int(kwargs["und_len"]) + packed_action_predictions = transformer.action_proj_out(hidden[indexes], per_noisy_domain_ids) + action = transformer._unpack_action_latents( + packed_action_predictions, + kwargs["action_token_shapes"], + kwargs["action_noisy_frame_indexes"], + ) + return _ModalityPredictions(vision=vision, sound=sound, action=action) + + def _onload_generation_modules(self) -> None: + if self._generation_modules_resident: + return + transformer = self._require_transformer() + for module in self._generation_modules(transformer): + module.to(self.device) + self._generation_modules_resident = True + + def _offload_generation_modules(self) -> None: + # A module.to() can fail partway through onload before the resident + # flag is set. Always visit the small generation heads during cleanup. + if self._transformer is None: + return + for module in self._generation_modules(self._transformer): + module.to(self.offload_device) + self._generation_modules_resident = False + + @staticmethod + def _generation_modules(transformer: torch.nn.Module) -> tuple[torch.nn.Module, ...]: + modules = ( + transformer.proj_in, + transformer.time_proj, + transformer.time_embedder, + transformer.norm_moe_gen, + transformer.proj_out, + ) + if hasattr(transformer, "audio_proj_in"): + modules += (transformer.audio_proj_in, transformer.audio_proj_out) + if hasattr(transformer, "action_proj_in"): + modules += (transformer.action_proj_in, transformer.action_proj_out) + return modules + + def _require_transformer(self) -> torch.nn.Module: + if self._transformer is None: + raise RuntimeError("runtime has not been bound to a transformer") + return self._transformer + + def _require_understanding_stack(self) -> UnderstandingStack: + if self._understanding_stack is None: + raise RuntimeError("runtime has not created the understanding stack") + return self._understanding_stack diff --git a/packages/diffusers-cpu-offloading/pyproject.toml b/packages/diffusers-cpu-offloading/pyproject.toml new file mode 100644 index 000000000..58797ca71 --- /dev/null +++ b/packages/diffusers-cpu-offloading/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools>=69"] +build-backend = "setuptools.build_meta" + +[project] +name = "cosmos3-streaming" +version = "0.1.0" +description = "External BF16 and FP8 CPU-offloaded execution for the Diffusers Cosmos 3 modular pipeline" +readme = "README.md" +requires-python = ">=3.12" +license = {text = "Apache-2.0"} +dependencies = [ + "accelerate>=0.31.0", + "av>=17.1.0", + "cosmos-guardrail>=0.3.1", + "diffusers @ git+https://github.com/huggingface/diffusers.git@main", + "imageio>=2.37.4", + "imageio-ffmpeg>=0.6.0", + # Cosmos Guardrail 0.3.1 reads NLTK data through Hugging Face cache symlinks. + "nltk==3.10.2", + "opencv-python>=4.10.0.84", + "sentencepiece>=0.1.91,!=0.1.92", + "torch>=2.7", + "transformers>=4.41.2", +] + +[project.optional-dependencies] +fp8 = ["nvidia-modelopt==0.44.0"] +benchmarks = ["scipy>=1.13", "cosmos-guardrail==0.3.1"] + +[tool.setuptools.packages.find] +include = ["cosmos3_streaming*"] + +[tool.ruff] +line-length = 120 +target-version = "py312" + +[tool.setuptools] +include-package-data = false From 29a37eca1c4ce29b5b96f3b7cd9127acd129ed9f Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Tue, 27 Oct 2026 15:46:13 -0700 Subject: [PATCH 2/2] Fix offloading package pre-commit formatting --- packages/diffusers-cpu-offloading/README.md | 19 +++++++++------ .../benchmarks/README.md | 23 +++++++++++-------- .../benchmarks/benchmark_bf16.py | 3 +++ .../benchmarks/benchmark_fp8.py | 3 +++ .../benchmarks/benchmark_safety.py | 3 +++ .../benchmarks/fp8_benchmark_audit.py | 3 +++ .../benchmarks/modalities.py | 3 +++ .../benchmarks/runner.py | 3 +++ .../cosmos3_streaming/__init__.py | 3 +++ .../cosmos3_streaming/attention/__init__.py | 3 +++ .../cosmos3_streaming/attention/base.py | 3 +++ .../cosmos3_streaming/attention/loader.py | 3 +++ .../cosmos3_streaming/attention/sdpa.py | 3 +++ .../cosmos3_streaming/block_views.py | 3 +++ .../cosmos3_streaming/config.py | 3 +++ .../cosmos3_streaming/fp8.py | 3 +++ .../cosmos3_streaming/memory_trace.py | 3 +++ .../cosmos3_streaming/metrics.py | 3 +++ .../cosmos3_streaming/modular.py | 3 +++ .../cosmos3_streaming/offloading.py | 3 +++ .../cosmos3_streaming/runtime.py | 3 +++ .../diffusers-cpu-offloading/pyproject.toml | 3 +++ 22 files changed, 86 insertions(+), 16 deletions(-) diff --git a/packages/diffusers-cpu-offloading/README.md b/packages/diffusers-cpu-offloading/README.md index 3b94cad16..ecd7786e5 100644 --- a/packages/diffusers-cpu-offloading/README.md +++ b/packages/diffusers-cpu-offloading/README.md @@ -1,3 +1,8 @@ + + # Cosmos3-Nano CPU-GPU weight streaming Run the Cosmos3-Nano Diffusers modular pipeline with model weights stored in @@ -104,13 +109,13 @@ on your hardware. Use the same `pipe` instance. Diffusers selects the generation workflow from the call arguments. -| Modality | Additional call inputs | -| --- | --- | -| Image-to-video | `image=` with a conditioning image | -| Video-to-video | `video=` and the appropriate frame-conditioning arguments | -| T2V or I2V with sound | `enable_sound=True` | -| Action | `action=CosmosActionCondition(...)` | -| Transfer | `control_videos=` with the matching control prompt and settings | +| Modality | Additional call inputs | +| --------------------- | --------------------------------------------------------------- | +| Image-to-video | `image=` with a conditioning image | +| Video-to-video | `video=` and the appropriate frame-conditioning arguments | +| T2V or I2V with sound | `enable_sound=True` | +| Action | `action=CosmosActionCondition(...)` | +| Transfer | `control_videos=` with the matching control prompt and settings | For example, add a conditioning image to the video call: diff --git a/packages/diffusers-cpu-offloading/benchmarks/README.md b/packages/diffusers-cpu-offloading/benchmarks/README.md index 1122b1579..c17017404 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/README.md +++ b/packages/diffusers-cpu-offloading/benchmarks/README.md @@ -1,3 +1,8 @@ + + # Reproduce Cosmos 3 Nano BF16 and FP8 benchmarks These runners generate complete examples using BF16 or the FP8 checkpoint's @@ -80,15 +85,15 @@ if you change the code or benchmark settings. ## What the matrix runs -| Modality flag | Example | Steps | Frames | FPS | Seed | Guidance | -|---|---|---:|---:|---:|---:|---:| -| `t2v` | Text to video | 35 | 189 | 24 | 123 | 6 | -| `i2v` | Image to video | 35 | 189 | 24 | 1111 | 6 | -| `t2vs` | Text to video with sound | 35 | 189 | 24 | 0 | 6 | -| `i2vs` | Image to video with sound | 35 | 189 | 24 | 0 | 6 | -| `action` | Action policy with video output | 30 | 17 | 5 | 0 | 1 | -| `v2v` | Video-conditioned generation | 35 | 189 | 24 | 0 | 6 | -| `transfer` | Edge-conditioned transfer | 35 per chunk | 121 | 30 | 0 | 3 | +| Modality flag | Example | Steps | Frames | FPS | Seed | Guidance | +| ------------- | ------------------------------- | -----------: | -----: | ---: | ---: | -------: | +| `t2v` | Text to video | 35 | 189 | 24 | 123 | 6 | +| `i2v` | Image to video | 35 | 189 | 24 | 1111 | 6 | +| `t2vs` | Text to video with sound | 35 | 189 | 24 | 0 | 6 | +| `i2vs` | Image to video with sound | 35 | 189 | 24 | 0 | 6 | +| `action` | Action policy with video output | 30 | 17 | 5 | 0 | 1 | +| `v2v` | Video-conditioned generation | 35 | 189 | 24 | 0 | 6 | +| `transfer` | Edge-conditioned transfer | 35 per chunk | 121 | 30 | 0 | 3 | Each modality runs at `240p` (432x240), `480p` (848x480), and `720p` (1280x720), in that order. Action uses its native resolution tiers 256, 480, diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py index e65baf08b..22842eefa 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_bf16.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Run full Cosmos3-Nano BF16 benchmarks with safety checking enabled by default.""" from runner import main diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py index 223f16f56..088d8f82e 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_fp8.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Full Cosmos3-Nano FP8 benchmark using the checkpoint's native Diffusers policy.""" from runner import main diff --git a/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py b/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py index 55915f496..670d8dce5 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py +++ b/packages/diffusers-cpu-offloading/benchmarks/benchmark_safety.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Shared, mandatory safety auditing for BF16 and FP8 benchmarks.""" import hashlib diff --git a/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py b/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py index 32fac4576..d53d9bd4b 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py +++ b/packages/diffusers-cpu-offloading/benchmarks/fp8_benchmark_audit.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Observe Diffusers' checkpoint-defined FP8 policy during full generation.""" from collections import Counter diff --git a/packages/diffusers-cpu-offloading/benchmarks/modalities.py b/packages/diffusers-cpu-offloading/benchmarks/modalities.py index e308eefdd..7a2018ebe 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/modalities.py +++ b/packages/diffusers-cpu-offloading/benchmarks/modalities.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import hashlib diff --git a/packages/diffusers-cpu-offloading/benchmarks/runner.py b/packages/diffusers-cpu-offloading/benchmarks/runner.py index 6c6f0bbf0..83718c4eb 100644 --- a/packages/diffusers-cpu-offloading/benchmarks/runner.py +++ b/packages/diffusers-cpu-offloading/benchmarks/runner.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Run full BF16 or checkpoint-policy FP8 benchmarks with safety checking enabled. Default: seven modalities, three resolutions, full steps, one GPU, fresh process diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py index 1141e9d57..cd43949a9 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/__init__.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Cosmos 3 Nano streaming execution for Diffusers modular pipelines.""" from .config import StreamingConfig diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py index 57e1bf81d..c1271274e 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/__init__.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from .base import CosmosAttentionBackend from .loader import load_attention_backend from .sdpa import TorchSdpaBackend diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py index 5add4421e..47b8e8b46 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/base.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from typing import Protocol, runtime_checkable diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py index 2f4fffd13..18d49e20f 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/loader.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from importlib import import_module diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py index b4c19f855..0c5970792 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/attention/sdpa.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import torch diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py index 31934a40c..5d2c6fdd5 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/block_views.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from collections.abc import Callable, Sequence diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py index 4d0ae438f..4dd92f206 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/config.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from dataclasses import dataclass, field diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py index 2da4efe9f..4189f955c 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/fp8.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from dataclasses import dataclass diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py index 55a323ab5..fe442a8ac 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/memory_trace.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import json diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py index be4c3fee2..d314b6362 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/metrics.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from dataclasses import dataclass diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py index 990890e50..ff7c537f3 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/modular.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations import inspect diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py index 46c3adcea..60b75c541 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/offloading.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + """Lifecycle operations for the Diffusers group-offload hooks we install.""" from __future__ import annotations diff --git a/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py b/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py index 4af34f21a..ab67853b4 100644 --- a/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py +++ b/packages/diffusers-cpu-offloading/cosmos3_streaming/runtime.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from __future__ import annotations from dataclasses import dataclass diff --git a/packages/diffusers-cpu-offloading/pyproject.toml b/packages/diffusers-cpu-offloading/pyproject.toml index 58797ca71..2be9ba425 100644 --- a/packages/diffusers-cpu-offloading/pyproject.toml +++ b/packages/diffusers-cpu-offloading/pyproject.toml @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + [build-system] requires = ["setuptools>=69"] build-backend = "setuptools.build_meta"