diff --git a/cookbooks/cosmos3/README.md b/cookbooks/cosmos3/README.md index 9dd21aa5..bd95381f 100644 --- a/cookbooks/cosmos3/README.md +++ b/cookbooks/cosmos3/README.md @@ -8,7 +8,7 @@ backend you want to run and follow that one section. | --- | --- | --- | | [Cosmos Framework](#cosmos-framework) | Native PyTorch inference, launched with `torchrun` | Reasoner, Generator (Audiovisual, Action, **Transfer**) | | [Diffusers](#diffusers) | Direct generation with `Cosmos3OmniPipeline` | Generator (Audiovisual) | -| [TensorRT-LLM Generator](#tensorrt-llm-generator) | OpenAI-compatible VisualGen server (image/video/audio/action/transfer generation) | Generator (Audiovisual, Action, **Transfer**) | +| [TensorRT-LLM Generator](#tensorrt-llm-generator) | Offline FP8 image/video generation and OpenAI-compatible VisualGen server (image/video/audio/action/transfer generation) | Generator (Audiovisual, Action, **Transfer**) | | [TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) | OpenAI-compatible image/video reasoning server | Reasoner | | [Transformers](#transformers) | Hugging Face Transformers inference | Reasoner | | [vLLM](#vllm) | OpenAI-compatible reasoning server (image/video understanding) | Reasoner | @@ -167,7 +167,7 @@ uv pip install --torch-backend=cu130 \ ## TensorRT-LLM Generator -OpenAI-compatible **VisualGen** server for Generator audiovisual text-to-image, +Offline FP8 checkpoint generation and an OpenAI-compatible **VisualGen** server for Generator audiovisual text-to-image, text-to-video, image-to-video, video-to-video, synchronized audio, Transfer, and Action examples. Initial Cosmos3 support was added in TensorRT-LLM PR @@ -229,7 +229,8 @@ For Python-only changes, the upstream guide also documents installing the checkout in editable mode. Then install the Cosmos3 guardrail package in the same environment unless you -explicitly disable guardrails before starting the server: +explicitly disable guardrails before running offline inference or starting the +server: ```bash pip install cosmos_guardrail==0.3.0 @@ -285,6 +286,40 @@ The separate `No safety models found, returning safe` warning in guardrail checks and face blurring remain configured; the NLTK workaround does not enable video-content classification or suppress that warning. +### Cosmos3 Nano and Super FP8 checkpoints (offline, single GPU) + +Cosmos3 Nano and Super publish ModelOpt-calibrated checkpoints on the `fp8` +revision of their Hugging Face repositories. These revisions contain the FP8 +weights, activation scales, and runtime policy. TensorRT-LLM reads that metadata +directly, so no quantization flag is required. Static-FP8 loading +([#17476](https://github.com/NVIDIA/TensorRT-LLM/pull/17476)) is merged on +TensorRT-LLM `main`; use a checkout at or after it rather than the older revision +pinned above. Select the FP8 weights with `--revision fp8`; TensorRT-LLM +downloads them on first use: + +```bash +python3 examples/visual_gen/models/cosmos3/cosmos3.py \ + --model nvidia/Cosmos3-Nano --revision fp8 \ + --visual_gen_args examples/visual_gen/configs/cosmos3-nano-1gpu.yaml \ + --prompt_file examples/visual_gen/models/cosmos3/prompts/t2v.json \ + --output_path nano_fp8_t2v.mp4 +``` + +Install `cosmos_guardrail==0.3.2` (not the `0.3.0` pin shown above) for these +offline runs; it needs no NLTK data setup. These FP8 checkpoints support one GPU +only, and the weights are about 21 GB for Nano and 70 GB for Super, so Super needs +a single GPU with enough memory for 70 GB of weights plus activations. Use the +BF16 checkpoints for tensor, Ulysses, context, or CFG parallelism, or for +parallel VAE. The +[FP8 checkpoint notebook](generator/audiovisual/run_fp8_with_trt_llm.ipynb) +runs Nano and Super through the offline TensorRT-LLM entry point for +text-to-image, text-to-video, image-to-video, and video-to-video generation, +text-to-video and image-to-video with synchronized audio, validates every PNG +and MP4 artifact, and serves the FP8 Reasoner for image and +video understanding. + +### VisualGen server + Set the TensorRT-LLM source root for the shared VisualGen config YAMLs. Run this from inside the TensorRT-LLM checkout — the directory the `git clone` above created, which is where `examples/` lives — or point `TRTLLM_ROOT` at that @@ -448,6 +483,28 @@ The server exposes `/health` and the OpenAI-compatible API at [Reasoner TensorRT-LLM notebook](reasoner/run_with_tensorrt_llm.ipynb) for image and video requests. +**Cosmos3-Nano or Cosmos3-Super FP8** (single GPU, port 8001). The `fp8` +revision carries ModelOpt-calibrated weights and static activation scales; +`trtllm-serve` reads them from the checkpoint, so no quantization flag is +needed. This needs a TensorRT-LLM build from `main`, as described in the +[FP8 checkpoint section](#cosmos3-nano-and-super-fp8-checkpoints-offline-single-gpu), +and Super FP8 fits on one Blackwell-class GPU: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +trtllm-serve nvidia/Cosmos3-Nano \ + --hf_revision fp8 \ + --host 0.0.0.0 \ + --port 8001 \ + --max_num_tokens 32768 +``` + +Replace `nvidia/Cosmos3-Nano` with `nvidia/Cosmos3-Super` for Super. On a Slurm +compute node where `trtllm-serve` stalls at `worker_initialization_wait`, +prefix the command with `mpirun -np 1 --bind-to none trtllm-llmapi-launch`. +The [FP8 checkpoint notebook](generator/audiovisual/run_fp8_with_trt_llm.ipynb) +starts this server and sends image and video requests to both models. + ## Transformers Local Python inference for the Cosmos3 Reasoner. This backend uses the diff --git a/cookbooks/cosmos3/generator/audiovisual/README.md b/cookbooks/cosmos3/generator/audiovisual/README.md index 627275a1..ae9ef3ce 100644 --- a/cookbooks/cosmos3/generator/audiovisual/README.md +++ b/cookbooks/cosmos3/generator/audiovisual/README.md @@ -421,6 +421,18 @@ students cover text-to-image and image-to-video only. Server launch commands live in the [shared environment setup guide](../../README.md#tensorrt-llm-generator). +### Cosmos3 Nano and Super FP8 notebook + +[`run_fp8_with_trt_llm.ipynb`](./run_fp8_with_trt_llm.ipynb) +runs the ModelOpt-calibrated Cosmos3 Nano and Super FP8 checkpoints (the `fp8` +revision on Hugging Face) directly through TensorRT-LLM's offline Cosmos3 entry point. It contains the complete +single-GPU commands for text-to-image, text-to-video, image-to-video, and +video-to-video generation, text-to-video and image-to-video with audio, on this folder's shared prompts and reference media, +followed by PNG and MP4 validation, plus image and video understanding through +the FP8 Reasoner served with `trtllm-serve`. Super needs a Blackwell-class single GPU. The checkpoint +metadata supplies the calibrated weight and activation scales; no runtime +quantization flag is required. + ## Run with NIM ### Quickstart diff --git a/cookbooks/cosmos3/generator/audiovisual/run_fp8_with_trt_llm.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_fp8_with_trt_llm.ipynb new file mode 100644 index 00000000..3a314dc2 --- /dev/null +++ b/cookbooks/cosmos3/generator/audiovisual/run_fp8_with_trt_llm.ipynb @@ -0,0 +1,834 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "markdown-00-95d2e8c3", + "metadata": {}, + "source": [ + "\n", + "\n", + "# Run Cosmos3 Nano and Super FP8 checkpoints with TensorRT-LLM\n", + "\n", + "TensorRT-LLM can run ModelOpt-calibrated Cosmos3 Nano and Super FP8\n", + "checkpoints directly. The checkpoint metadata supplies the calibrated weight\n", + "and activation scales, so no quantization CLI flag is\n", + "needed.\n", + "\n", + "This notebook covers the single-GPU generation modes, with and without audio,\n", + "and the Reasoner:\n", + "\n", + "| Mode | Prompt or condition | Output |\n", + "| -------------------- | ------------------------------------- | --------- |\n", + "| Text to image (T2I) | Structured text prompt | PNG image |\n", + "| Text to video (T2V) | Structured text prompt | MP4 video |\n", + "| Image to video (I2V) | Structured prompt and reference image | MP4 video |\n", + "| Video to video (V2V) | Structured prompt and reference video | MP4 video |\n", + "| Text to video with audio (T2AV) | Structured prompt | MP4 video with stereo audio |\n", + "| Image to video with audio (TI2AV) | Structured prompt and reference image | MP4 video with stereo audio |\n", + "| Reasoner | Image or video plus a question | Text |\n", + "\n", + "FP8 Cosmos3 is single-GPU only. TensorRT-LLM rejects tensor parallelism,\n", + "Ulysses parallelism, context parallelism, CFG parallelism, or parallel VAE\n", + "sizes greater than one for these checkpoints. Use a BF16 checkpoint for a\n", + "multi-GPU deployment.\n", + "\n", + "## 1. Prerequisites\n", + "\n", + "Build and install TensorRT-LLM from its `main` branch by following the\n", + "[TensorRT-LLM source-build guide](https://nvidia.github.io/TensorRT-LLM/installation/build-from-source.html).\n", + "Static-FP8 Cosmos3 loading\n", + "([#17476](https://github.com/NVIDIA/TensorRT-LLM/pull/17476)) is merged on\n", + "`main`, so a checkout at or after that change is required. The source checkout\n", + "supplies the Cosmos3 example script and one-GPU configuration files used\n", + "below. Prompts, negative prompts, the reference image, and the reference video\n", + "come from this cookbook's shared audiovisual assets, so the same inputs run\n", + "through every other backend in this folder.\n", + "\n", + "Install `ffmpeg` so TensorRT-LLM can encode MP4 output. Install the Cosmos3\n", + "guardrail package, accept the gated\n", + "[`nvidia/Cosmos-1.0-Guardrail`](https://huggingface.co/nvidia/Cosmos-1.0-Guardrail)\n", + "terms, and authenticate to Hugging Face before generation:\n", + "\n", + "```shell\n", + "apt-get update && apt-get install -y ffmpeg\n", + "pip install cosmos_guardrail==0.3.2\n", + "pip uninstall -y opencv-python\n", + "pip install opencv-python-headless\n", + "hf auth login\n", + "```\n", + "\n", + "TensorRT-LLM downloads only the guardrail files it loads, and no NLTK data\n", + "workaround is needed. The guardrail package performs prompt text checks and\n", + "RetinaFace face blurring. Its generated-frame content classifier is disabled\n", + "upstream because it produced too many false positives.\n", + "\n", + "Nano and Super publish their ModelOpt FP8 checkpoints on the `fp8` revision of\n", + "their Hugging Face repositories, next to the BF16 weights. The commands below\n", + "pass `--revision fp8` so TensorRT-LLM downloads that revision on first use. The\n", + "weights are about 21 GB for Nano and 70 GB for Super, and the weights cannot be\n", + "split across GPUs, so Super needs a single GPU with enough memory for 70 GB of\n", + "weights plus activations, such as a B200 or B300. To use a local copy instead,\n", + "`hf download` the `fp8` revision and point `COSMOS3_NANO_FP8` or\n", + "`COSMOS3_SUPER_FP8` at that directory; `--revision` is ignored for local paths.\n", + "\n", + "Set the paths used by the commands in this notebook. You may export these\n", + "variables before starting Jupyter or replace the placeholder defaults below:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "python-01-bdb74719", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import os\n", + "\n", + "\n", + "def find_repo_root(start: Path) -> Path:\n", + " for path in [start, *start.parents]:\n", + " if (path / \"README.md\").exists() and (path / \"cookbooks\").exists():\n", + " return path\n", + " return start\n", + "\n", + "\n", + "COSMOS_ROOT = find_repo_root(Path.cwd().resolve())\n", + "ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"generator\" / \"audiovisual\" / \"assets\"\n", + "COSMOS3_T2I_PROMPT = ASSETS / \"prompts\" / \"text2image\" / \"robot_draping.json\"\n", + "COSMOS3_T2V_PROMPT = ASSETS / \"prompts\" / \"text2video\" / \"robot_kitchen.json\"\n", + "COSMOS3_T2V_NEGATIVE = ASSETS / \"negative_prompts\" / \"text2video\" / \"neg_prompt.json\"\n", + "COSMOS3_I2V_PROMPT = ASSETS / \"prompts\" / \"image2video\" / \"car_driving.json\"\n", + "COSMOS3_I2V_NEGATIVE = ASSETS / \"negative_prompts\" / \"image2video\" / \"neg_prompt.json\"\n", + "COSMOS3_I2V_IMAGE = ASSETS / \"images\" / \"image2video\" / \"car_driving.jpg\"\n", + "COSMOS3_V2V_VIDEO = ASSETS / \"videos\" / \"car_driving_plain.mp4\"\n", + "COSMOS3_T2AV_PROMPT = ASSETS / \"prompts\" / \"text2video\" / \"robot_pouring_water_audio.json\"\n", + "COSMOS3_TI2AV_PROMPT = ASSETS / \"prompts\" / \"image2video\" / \"coastal_road_audio.json\"\n", + "COSMOS3_TI2AV_IMAGE = ASSETS / \"images\" / \"image2video\" / \"coastal_road_audio.jpg\"\n", + "\n", + "TRTLLM_ROOT = Path(os.environ.get(\"TRTLLM_ROOT\", \"/path/to/TensorRT-LLM\")).expanduser().resolve()\n", + "# A Hugging Face repo ID (loaded at --revision fp8) or a local copy of the fp8 revision.\n", + "COSMOS3_NANO_FP8 = os.environ.get(\"COSMOS3_NANO_FP8\", \"nvidia/Cosmos3-Nano\")\n", + "COSMOS3_SUPER_FP8 = os.environ.get(\"COSMOS3_SUPER_FP8\", \"nvidia/Cosmos3-Super\")\n", + "COSMOS3_FP8_OUTPUTS = Path(\n", + " os.environ.get(\"COSMOS3_FP8_OUTPUTS\", Path.cwd() / \"outputs\" / \"tensorrt_llm_fp8\")\n", + ").expanduser().resolve()\n", + "\n", + "inputs = {\n", + " \"TRTLLM_ROOT\": TRTLLM_ROOT,\n", + " \"COSMOS3_T2I_PROMPT\": COSMOS3_T2I_PROMPT,\n", + " \"COSMOS3_T2V_PROMPT\": COSMOS3_T2V_PROMPT,\n", + " \"COSMOS3_T2V_NEGATIVE\": COSMOS3_T2V_NEGATIVE,\n", + " \"COSMOS3_I2V_PROMPT\": COSMOS3_I2V_PROMPT,\n", + " \"COSMOS3_I2V_NEGATIVE\": COSMOS3_I2V_NEGATIVE,\n", + " \"COSMOS3_I2V_IMAGE\": COSMOS3_I2V_IMAGE,\n", + " \"COSMOS3_V2V_VIDEO\": COSMOS3_V2V_VIDEO,\n", + " \"COSMOS3_T2AV_PROMPT\": COSMOS3_T2AV_PROMPT,\n", + " \"COSMOS3_TI2AV_PROMPT\": COSMOS3_TI2AV_PROMPT,\n", + " \"COSMOS3_TI2AV_IMAGE\": COSMOS3_TI2AV_IMAGE,\n", + "}\n", + "for name, path in inputs.items():\n", + " if not path.exists():\n", + " raise FileNotFoundError(f\"Set {name} to an existing path; got {path}\")\n", + "\n", + "COSMOS3_FP8_OUTPUTS.mkdir(parents=True, exist_ok=True)\n", + "for name, value in {\n", + " **inputs,\n", + " \"COSMOS3_NANO_FP8\": COSMOS3_NANO_FP8,\n", + " \"COSMOS3_SUPER_FP8\": COSMOS3_SUPER_FP8,\n", + " \"COSMOS3_FP8_OUTPUTS\": COSMOS3_FP8_OUTPUTS,\n", + "}.items():\n", + " os.environ[name] = str(value)\n", + " print(f\"{name}={value}\")" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-02-5954b927", + "metadata": {}, + "source": [ + "\n", + "The checkpoints may carry a `diffusion_step_policy` in their quantization\n", + "metadata. TensorRT-LLM applies that checkpoint-owned policy automatically; do\n", + "not add a separate runtime override for it.\n", + "\n", + "## 2. Cosmos3 Nano FP8\n", + "\n", + "### Text to image\n", + "\n", + "The text-to-image config warms the 1024x1024, one-frame shape. The cookbook's\n", + "structured T2I prompt matches the checkpoint's trained prompt format, and\n", + "`--output_type image` makes the expected output explicit. `--revision fp8`\n", + "selects the FP8 weights.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-03-0f16ab37", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2I_PROMPT\" \\\n", + " --output_type image \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_t2i.png\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-04-36f6f594", + "metadata": {}, + "source": [ + "\n", + "### Text to video\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-05-ca0df107", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_T2V_NEGATIVE\" \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_t2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-06-4f217480", + "metadata": {}, + "source": [ + "\n", + "### Image to video\n", + "\n", + "`--image_path` supplies the reference image as a local path, `file://` URL,\n", + "HTTP(S) URL, or `data:` URI.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-07-baadc896", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_I2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_I2V_NEGATIVE\" \\\n", + " --image_path \"$COSMOS3_I2V_IMAGE\" \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_i2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-08-9c92db9f", + "metadata": {}, + "source": [ + "\n", + "### Video to video\n", + "\n", + "This command conditions on the cookbook's reference driving clip. Latent\n", + "indexes `[0, 1]` with `\"first\"` keep the first five decoded pixel frames as\n", + "clean conditioning, matching the other backends in this folder.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-09-ab023113", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_I2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_I2V_NEGATIVE\" \\\n", + " --use_system_prompt \\\n", + " --video_path \"$COSMOS3_V2V_VIDEO\" \\\n", + " --extra_params '{\"condition_video_latent_indexes\": [0, 1], \"condition_video_keep\": \"first\"}' \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_v2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-09a-t2av", + "metadata": {}, + "source": [ + "\n", + "### Text to video with audio\n", + "\n", + "`--enable_audio` generates a synchronized stereo audio track with the video.\n", + "The FP8 checkpoints carry the audio tower, so no other change is needed.\n" + ] + }, + { + "cell_type": "code", + "id": "bash-09a-t2av", + "metadata": {}, + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2AV_PROMPT\" \\\n", + " --enable_audio \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_t2av.mp4\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-09a-ti2av", + "metadata": {}, + "source": [ + "\n", + "### Image to video with audio\n" + ] + }, + { + "cell_type": "code", + "id": "bash-09a-ti2av", + "metadata": {}, + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_NANO_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_TI2AV_PROMPT\" \\\n", + " --image_path \"$COSMOS3_TI2AV_IMAGE\" \\\n", + " --enable_audio \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/nano_ti2av.mp4\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-10-e177f899", + "metadata": {}, + "source": [ + "\n", + "## 3. Cosmos3 Super FP8\n", + "\n", + "The same one-GPU interfaces apply to the Super checkpoint. Do not use\n", + "`cosmos3-super-4gpu.yaml`; these FP8 checkpoints reject that parallel\n", + "configuration.\n", + "\n", + "### Text to image\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-11-15a5a2ad", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2I_PROMPT\" \\\n", + " --output_type image \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_t2i.png\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-12-eb633740", + "metadata": {}, + "source": [ + "\n", + "### Text to video\n", + "\n", + "The file name `cosmos3-nano-1gpu.yaml` is historical: TensorRT-LLM documents\n", + "the configuration as the shared one-GPU config for both Nano and Super.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-13-7601d3d3", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_T2V_NEGATIVE\" \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_t2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-14-8055a556", + "metadata": {}, + "source": [ + "\n", + "### Image to video\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-15-0056f0db", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_I2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_I2V_NEGATIVE\" \\\n", + " --image_path \"$COSMOS3_I2V_IMAGE\" \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_i2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-16-b92c6894", + "metadata": {}, + "source": [ + "\n", + "### Video to video\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-17-0579fa3f", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_I2V_PROMPT\" \\\n", + " --negative_prompt_file \"$COSMOS3_I2V_NEGATIVE\" \\\n", + " --use_system_prompt \\\n", + " --video_path \"$COSMOS3_V2V_VIDEO\" \\\n", + " --extra_params '{\"condition_video_latent_indexes\": [0, 1], \"condition_video_keep\": \"first\"}' \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_v2v.mp4\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-17a-t2av", + "metadata": {}, + "source": [ + "\n", + "### Text to video with audio\n", + "\n", + "`--enable_audio` generates a synchronized stereo audio track with the video.\n", + "The FP8 checkpoints carry the audio tower, so no other change is needed.\n" + ] + }, + { + "cell_type": "code", + "id": "bash-17a-t2av", + "metadata": {}, + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_T2AV_PROMPT\" \\\n", + " --enable_audio \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_t2av.mp4\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-17a-ti2av", + "metadata": {}, + "source": [ + "\n", + "### Image to video with audio\n" + ] + }, + { + "cell_type": "code", + "id": "bash-17a-ti2av", + "metadata": {}, + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python \"$TRTLLM_ROOT/examples/visual_gen/models/cosmos3/cosmos3.py\" \\\n", + " --model \"$COSMOS3_SUPER_FP8\" \\\n", + " --revision fp8 \\\n", + " --visual_gen_args \"$TRTLLM_ROOT/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml\" \\\n", + " --prompt_file \"$COSMOS3_TI2AV_PROMPT\" \\\n", + " --image_path \"$COSMOS3_TI2AV_IMAGE\" \\\n", + " --enable_audio \\\n", + " --output_path \"$COSMOS3_FP8_OUTPUTS/super_ti2av.mp4\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-18-674ba7b9", + "metadata": {}, + "source": [ + "\n", + "## 4. Verify the artifacts\n", + "\n", + "Decode the two PNG files and print their dimensions:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-19-41172619", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "python - \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_t2i.png\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_t2i.png\" <<'PY'\n", + "from PIL import Image\n", + "import sys\n", + "\n", + "for path in sys.argv[1:]:\n", + " with Image.open(path) as image:\n", + " image.verify()\n", + " print(path, image.format, image.size)\n", + "PY\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-20-39accea0", + "metadata": {}, + "source": [ + "\n", + "Inspect and fully decode every generated MP4, and confirm that the audio\n", + "outputs carry an audio stream:\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bash-21-b400c1f9", + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "set -euo pipefail\n", + "\n", + "for output in \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_t2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_i2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_v2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_t2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_i2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_v2v.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_t2av.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/nano_ti2av.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_t2av.mp4\" \\\n", + " \"$COSMOS3_FP8_OUTPUTS/super_ti2av.mp4\"; do\n", + " ffprobe -v error \\\n", + " -show_entries stream=codec_type,codec_name,width,height,r_frame_rate,nb_frames \\\n", + " -of json \\\n", + " \"$output\"\n", + " ffmpeg -v error -i \"$output\" -f null -\n", + "done\n", + "\n", + "# The audio outputs must also carry a decodable audio stream.\n", + "for output in \"$COSMOS3_FP8_OUTPUTS\"/*_t2av.mp4 \"$COSMOS3_FP8_OUTPUTS\"/*_ti2av.mp4; do\n", + " audio=$(ffprobe -v error -select_streams a:0 \\\n", + " -show_entries stream=codec_name,channels,sample_rate,duration -of csv=p=0 \"$output\")\n", + " [ -n \"$audio\" ] || { echo \"no audio stream in $output\"; exit 1; }\n", + " echo \"$(basename \"$output\") audio: $audio\"\n", + "done\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-22-334c3588", + "metadata": {}, + "source": [ + "\n", + "A successful encode/decode proves the interface and artifact contract; it is\n", + "not a claim that FP8 output quality matches BF16.\n" + ] + }, + { + "cell_type": "markdown", + "id": "markdown-23-reasoner-intro", + "metadata": {}, + "source": [ + "## 5. Reasoner: image and video understanding\n", + "\n", + "The same `fp8` revision also serves the Cosmos3 Reasoner through `trtllm-serve`,\n", + "TensorRT-LLM's OpenAI-compatible LLM server. It loads the language and vision\n", + "towers only: the text-layer linear layers run in FP8 with the checkpoint's\n", + "static scales, while the vision tower, embeddings, and output head stay in\n", + "BF16. The generator's diffusion step policy does not apply to this path.\n", + "\n", + "The helpers below start `trtllm-serve` in the background with\n", + "`--hf_revision fp8`, wait for `/health`, send image and video requests with the\n", + "`openai` client, and stop the server. The server listens on port 8001, the same\n", + "endpoint as the [Reasoner notebook](../../reasoner/run_with_tensorrt_llm.ipynb),\n", + "so every example there also runs against it. The first start downloads the\n", + "weights and can take several minutes.\n", + "\n", + "If `trtllm-serve` stalls at `worker_initialization_wait` on a Slurm compute\n", + "node, MPI cannot spawn its worker process there. Set\n", + "`COSMOS3_TRTLLM_LAUNCHER=\"mpirun -np 1 --bind-to none trtllm-llmapi-launch\"`\n", + "before starting Jupyter to launch the server under `mpirun` instead.\n" + ] + }, + { + "cell_type": "code", + "id": "python-24-reasoner-helpers", + "metadata": {}, + "source": [ + "import shlex\n", + "import subprocess\n", + "import time\n", + "import urllib.request\n", + "\n", + "import openai\n", + "\n", + "REASONER_ASSETS = COSMOS_ROOT / \"cookbooks\" / \"cosmos3\" / \"reasoner\" / \"assets\"\n", + "REASONER_PORT = 8001\n", + "REASONER_URL = f\"http://localhost:{REASONER_PORT}\"\n", + "VIDEO_KWARGS = {\n", + " \"media_io_kwargs\": {\"video\": {\"num_frames\": -1, \"fps\": 4}},\n", + " \"mm_processor_kwargs\": {\"do_sample_frames\": False},\n", + "}\n", + "# Optional launcher prefix. Leave unset on a workstation or in a container. On a\n", + "# Slurm node where MPI cannot spawn worker processes, set for example:\n", + "# COSMOS3_TRTLLM_LAUNCHER=\"mpirun -np 1 --bind-to none trtllm-llmapi-launch\"\n", + "REASONER_LAUNCHER = shlex.split(os.environ.get(\"COSMOS3_TRTLLM_LAUNCHER\", \"\"))\n", + "reasoner_server = None\n", + "\n", + "\n", + "def start_reasoner(model: str, timeout_s: int = 3600) -> None:\n", + " \"\"\"Start trtllm-serve for the FP8 Reasoner and wait until /health answers.\"\"\"\n", + " global reasoner_server\n", + " command = REASONER_LAUNCHER + [\n", + " \"trtllm-serve\", model,\n", + " \"--host\", \"127.0.0.1\", \"--port\", str(REASONER_PORT),\n", + " \"--max_num_tokens\", \"32768\",\n", + " ]\n", + " if not Path(model).exists():\n", + " command += [\"--hf_revision\", \"fp8\"] # Hub repo ID: select the FP8 branch\n", + " log_path = COSMOS3_FP8_OUTPUTS / \"reasoner_server.log\"\n", + " reasoner_server = subprocess.Popen(\n", + " command, stdout=log_path.open(\"w\"), stderr=subprocess.STDOUT\n", + " )\n", + " print(\" \".join(command))\n", + " deadline = time.time() + timeout_s\n", + " while time.time() < deadline:\n", + " if reasoner_server.poll() is not None:\n", + " raise RuntimeError(f\"trtllm-serve exited early; see {log_path}\")\n", + " try:\n", + " with urllib.request.urlopen(f\"{REASONER_URL}/health\", timeout=5):\n", + " print(\"Reasoner server ready\")\n", + " return\n", + " except OSError:\n", + " time.sleep(5)\n", + " stop_reasoner()\n", + " raise TimeoutError(f\"trtllm-serve was not ready after {timeout_s} s; see {log_path}\")\n", + "\n", + "\n", + "def stop_reasoner() -> None:\n", + " if reasoner_server is not None and reasoner_server.poll() is None:\n", + " reasoner_server.terminate()\n", + " try:\n", + " reasoner_server.wait(timeout=120)\n", + " except subprocess.TimeoutExpired:\n", + " reasoner_server.kill()\n", + "\n", + "\n", + "def ask(media: str, prompt: str) -> str:\n", + " client = openai.OpenAI(api_key=\"EMPTY\", base_url=f\"{REASONER_URL}/v1\")\n", + " kind = \"video_url\" if media.endswith(\".mp4\") else \"image_url\"\n", + " response = client.chat.completions.create(\n", + " model=client.models.list().data[0].id,\n", + " messages=[{\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": kind, kind: {\"url\": (REASONER_ASSETS / media).resolve().as_uri()}},\n", + " {\"type\": \"text\", \"text\": prompt},\n", + " ],\n", + " }],\n", + " max_tokens=4096,\n", + " seed=0,\n", + " extra_body=VIDEO_KWARGS if kind == \"video_url\" else {},\n", + " )\n", + " answer = response.choices[0].message.content\n", + " assert answer and answer.strip(), f\"empty answer for {media}\"\n", + " return answer\n", + "\n", + "\n", + "REASONER_EXAMPLES = [\n", + " (\"robot_153.jpg\", \"Caption the image in detail.\"),\n", + " (\"video_caption.mp4\", \"Describe the video in detail.\"),\n", + " (\"grounding_2d.png\", \"Locate the accurate bounding box of the load as a whole. Return a json.\"),\n", + "]\n", + "\n", + "\n", + "def run_reasoner_examples() -> None:\n", + " for media, prompt in REASONER_EXAMPLES:\n", + " print(f\"### {media}: {prompt}\")\n", + " print(ask(media, prompt), end=\"\\n\\n\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-25-reasoner-nano", + "metadata": {}, + "source": [ + "### Cosmos3 Nano FP8 Reasoner\n" + ] + }, + { + "cell_type": "code", + "id": "python-26-reasoner-nano", + "metadata": {}, + "source": [ + "start_reasoner(COSMOS3_NANO_FP8)\n", + "try:\n", + " run_reasoner_examples()\n", + "finally:\n", + " stop_reasoner()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-27-reasoner-super", + "metadata": {}, + "source": [ + "### Cosmos3 Super FP8 Reasoner\n", + "\n", + "Super's FP8 language model fits one Blackwell-class GPU; no tensor parallelism\n", + "is needed.\n" + ] + }, + { + "cell_type": "code", + "id": "python-28-reasoner-super", + "metadata": {}, + "source": [ + "start_reasoner(COSMOS3_SUPER_FP8)\n", + "try:\n", + " run_reasoner_examples()\n", + "finally:\n", + " stop_reasoner()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "markdown-29-scope", + "metadata": {}, + "source": [ + "## 6. Scope\n", + "\n", + "The FP8 checkpoints run on one GPU only, and no FP8-versus-BF16 quality\n", + "benchmark, for video or audio, is claimed here.\n", + "\n", + "The Reasoner section exercises three representative requests per model; the\n", + "[Reasoner notebook](../../reasoner/run_with_tensorrt_llm.ipynb) has the full set\n", + "of image and video understanding examples, which run unchanged against the FP8\n", + "server.\n", + "\n", + "For the complete TensorRT-LLM Cosmos3 example contract, including serving,\n", + "guardrails, media dependencies, and other checkpoint families, see the\n", + "[TensorRT-LLM Cosmos3 example](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/visual_gen/models/cosmos3).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}