From 39b17e0ea38200ebfa729c9a3cc0589fdc12091a Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Wed, 5 Aug 2026 14:02:36 -0700 Subject: [PATCH 1/7] cookbook(diffusers): add FP8 quantized checkpoint loading path Enable Nano/Super and Super 4-Step FP8 via COSMOS3_DIFFUSERS_PRECISION without duplicating request cells. Load public MODEL_IDS with revision=fp8 or a local COSMOS3_FP8_MODEL_PATH. --- .../audiovisual/run_with_diffusers.ipynb | 115 ++++++++++++++++-- 1 file changed, 104 insertions(+), 11 deletions(-) diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 2c058af6..320aa015 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -19,6 +19,8 @@ "\n", "Run the Cosmos3-Nano, Cosmos3-Super, or Cosmos3-Edge examples independently. Cosmos3-Edge has no audio modules and uses its 480p generation settings. Run Cosmos3-Super T2V/I2V examples without audio. Each section loads the matching model explicitly.\n", "\n", + "**FP8:** `Cosmos3-Nano`, `Cosmos3-Super`, and the Super 4-Step distilled checkpoints also ship static-scale FP8 builds. Enable them with `COSMOS3_DIFFUSERS_PRECISION=fp8` (see **Quantized Checkpoints** below). Request cells stay the same; only checkpoint loading changes. `Cosmos3-Edge` has no FP8 build.\n", + "\n", "Note: if you have already completed steps 1-3 and installed the `Cosmos3 Diffusers (Python 3.13)` kernel, switch to that kernel and jump directly to step 4. Run the restore cell there, then continue with verification and the examples.\n" ] }, @@ -57,6 +59,8 @@ "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", + "# Optional: bf16 (default) or fp8 for Nano/Super/4-Step — see Quantized Checkpoints\n", + "export COSMOS3_DIFFUSERS_PRECISION=bf16\n", "```\n" ] }, @@ -103,6 +107,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -114,6 +119,7 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -159,7 +165,8 @@ " ipykernel \\\n", " torch \\\n", " torchvision \\\n", - " transformers\n", + " transformers \\\n", + " \"nvidia-modelopt==0.44.0\"\n", "\n", "\"$COSMOS3_DIFFUSERS_VENV/bin/python\" -m ipykernel install --user \\\n", " --name cosmos3-diffusers \\\n", @@ -631,10 +638,38 @@ " \"Cosmos3-Super-Image2Video-4Step\",\n", "}\n", "\n", + "# FP8 builds for Nano/Super and Super 4-Step (same request path as bf16; see Quantized Checkpoints).\n", + "FP8_SUPPORTED_MODELS = {\n", + " \"Cosmos3-Nano\",\n", + " \"Cosmos3-Super\",\n", + " \"Cosmos3-Super-Text2Image-4Step\",\n", + " \"Cosmos3-Super-Image2Video-4Step\",\n", + "}\n", + "\n", "_pipe = None\n", "_pipe_model = None\n", "\n", "\n", + "def precision_mode() -> str:\n", + " return os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\").strip().lower()\n", + "\n", + "\n", + "def is_fp8_precision() -> bool:\n", + " return precision_mode() == \"fp8\"\n", + "\n", + "\n", + "def resolve_fp8_pretrained_source(model: str):\n", + " \"\"\"Return a local path or public HF repo id for the FP8 checkpoint.\n", + "\n", + " Uses `MODEL_IDS[model]` with `revision=COSMOS3_FP8_REVISION` (default `fp8`) unless\n", + " `COSMOS3_FP8_MODEL_PATH` points at a local Diffusers checkpoint directory.\n", + " \"\"\"\n", + " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " if override:\n", + " return Path(override).expanduser().resolve()\n", + " return resolve_model_id(model)\n", + "\n", + "\n", "def resolve_model_id(model: str) -> str:\n", " return MODEL_IDS.get(model, model)\n", "\n", @@ -680,12 +715,35 @@ "def get_pipe(model: str):\n", " global _pipe, _pipe_model\n", " model_id = resolve_model_id(model)\n", - " if _pipe is not None and _pipe_model == model_id:\n", + " prec = precision_mode()\n", + " cache_key = f\"{model_id}@{prec}\"\n", + " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", - " print(f\"loading {model_id}...\")\n", " t0 = time.time()\n", + " fp8 = is_fp8_precision()\n", + " if fp8 and model not in FP8_SUPPORTED_MODELS:\n", + " raise RuntimeError(\n", + " f\"FP8 is only supported for {sorted(FP8_SUPPORTED_MODELS)}; got {model}. \"\n", + " \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n", + " )\n", + " if fp8:\n", + " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", + " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", + " from modelopt.torch.opt import enable_huggingface_checkpointing\n", + "\n", + " enable_huggingface_checkpointing()\n", + " pretrained_source = resolve_fp8_pretrained_source(model)\n", + " from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", + " if not os.environ.get(\"COSMOS3_FP8_MODEL_PATH\"):\n", + " from_kwargs[\"revision\"] = os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\")\n", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " pretrained_source = model_id\n", + " from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", + " print(f\"loading {model_id}...\")\n", + "\n", " if is_distilled_model(model):\n", " if Cosmos3DistilledModularPipeline is None:\n", " raise RuntimeError(\n", @@ -693,24 +751,24 @@ " \"Reinstall diffusers from a revision that includes the Cosmos3 distilled pipeline.\"\n", " )\n", " pipe = Cosmos3DistilledModularPipeline.from_pretrained(\n", - " model_id,\n", - " token=os.environ.get(\"HF_TOKEN\") or None,\n", + " pretrained_source,\n", + " **from_kwargs,\n", " )\n", " pipe.load_components(torch_dtype=torch.bfloat16)\n", " pipe.enable_safety_checker()\n", " else:\n", " pipe = Cosmos3OmniPipeline.from_pretrained(\n", - " model_id,\n", + " pretrained_source,\n", " torch_dtype=torch.bfloat16,\n", " safety_checker=None,\n", " enable_safety_checker=True,\n", - " token=os.environ.get(\"HF_TOKEN\") or None,\n", + " **from_kwargs,\n", " )\n", " pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FIXED_SAMPLING[\"shift\"])\n", " pipe.to(\"cuda\")\n", " _pipe = pipe\n", - " _pipe_model = model_id\n", - " print(f\"loaded pipeline in {time.time() - t0:.1f}s\")\n", + " _pipe_model = cache_key\n", + " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec})\")\n", " return _pipe\n", "\n", "\n", @@ -839,10 +897,45 @@ " for src in images:\n", " print(f\"source: {src} ({src.stat().st_size // 1024} KB)\")\n", " display(Image(filename=str(src), width=720))\n" - ], + ] + }, + { + "cell_type": "markdown", + "id": "quantized-checkpoints-fp8", + "metadata": {}, + "source": [ + "## Quantized Checkpoints (FP8)\n", + "\n", + "FP8 builds are available for **`Cosmos3-Nano`**, **`Cosmos3-Super`**, and the Super **4-Step** distilled checkpoints (`Cosmos3-Edge` has no FP8 build). Serving FP8 is a checkpoint-loading change: the use-case request cells below stay the same.\n", + "\n", + "Static-scale FP8 checkpoints are produced with [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer) and load through the same Diffusers pipelines as bf16 (`Cosmos3OmniPipeline` for Nano/Super, `Cosmos3DistilledModularPipeline` for 4-Step). The install cell pins `nvidia-modelopt==0.44.0`.\n", + "\n", + "Enable FP8 before running the helper cell (or re-run the helper cell after changing the env):\n", + "\n", + "```bash\n", + "export COSMOS3_DIFFUSERS_PRECISION=fp8\n", + "```\n", + "\n", + "With FP8 enabled, this notebook loads the public model id from `MODEL_IDS` with `revision=fp8` (override with `COSMOS3_FP8_REVISION`). Optionally set `COSMOS3_FP8_MODEL_PATH` to a local Diffusers checkpoint directory instead.\n", + "\n", + "Loading may print a `Some weights of the model checkpoint ... not used` warning for `weight_scale` / `input_scale` keys. Those are vLLM-Omni scale tensors shipped alongside the Diffusers ModelOpt state; the warning is expected.\n" + ] + }, + { + "cell_type": "code", "execution_count": null, + "id": "quantized-checkpoints-fp8-enable", + "metadata": {}, "outputs": [], - "id": "77e18231" + "source": [ + "# Optional: flip precision for subsequent get_pipe / run_diffusers_payload calls.\n", + "# Re-run the helper cell above after changing this if helpers were already executed.\n", + "import os\n", + "\n", + "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\" # uncomment to use FP8 Nano/Super/4-Step\n", + "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n", + "print(\"COSMOS3_FP8_REVISION:\", os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\"))\n" + ] }, { "cell_type": "markdown", From f4878b6c001f12338650f04003393939c625e603 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Thu, 6 Aug 2026 11:29:11 -0700 Subject: [PATCH 2/7] cookbook(diffusers): fix Distilled FP8 ModelOpt restore via transformer preload Modular load_components passes repo+subfolder, which ModelOpt 0.44 ignores when looking up modelopt_state.pth. Pre-load transformer from its full local path, pin it, then load remaining components. Add verify_fp8 helper. --- .../audiovisual/run_with_diffusers.ipynb | 61 ++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 320aa015..d83734e9 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -678,6 +678,43 @@ " return model in DISTILLED_MODELS\n", "\n", "\n", + "def materialize_checkpoint_dir(pretrained_source, from_kwargs: dict) -> Path:\n", + " \"\"\"Return a local checkpoint root (download from Hub when given a repo id).\"\"\"\n", + " path = Path(str(pretrained_source))\n", + " if path.is_dir():\n", + " return path.resolve()\n", + " from huggingface_hub import snapshot_download\n", + "\n", + " return Path(\n", + " snapshot_download(\n", + " str(pretrained_source),\n", + " revision=from_kwargs.get(\"revision\"),\n", + " token=from_kwargs.get(\"token\"),\n", + " )\n", + " ).resolve()\n", + "\n", + "\n", + "def distilled_components_root(checkpoint_dir: Path) -> str:\n", + " \"\"\"Local Diffusers root for modular `load_components` after transformer preload.\"\"\"\n", + " return str(Path(checkpoint_dir).resolve())\n", + "\n", + "\n", + "def verify_fp8(model: str) -> None:\n", + " \"\"\"Confirm ModelOpt restored FP8 weights and quantizers (not just float8 tensors).\"\"\"\n", + " transformer = get_pipe(model).transformer\n", + " tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n", + "\n", + " fp8 = sum(p.dtype == torch.float8_e4m3fn for p in transformer.parameters())\n", + " quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n", + " meta = sum(t.is_meta for _, t in tensors)\n", + "\n", + " # Float8 weights alone are insufficient; quantizers prove ModelOpt restored the scales.\n", + " if not fp8 or not quantizers or meta:\n", + " raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + "\n", + " print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta=0\")\n", + "\n", + "\n", "def cuda_allocated_gib() -> float:\n", " return torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0\n", "\n", @@ -754,7 +791,27 @@ " pretrained_source,\n", " **from_kwargs,\n", " )\n", - " pipe.load_components(torch_dtype=torch.bfloat16)\n", + " if fp8:\n", + " # Pre-load the transformer from its full path so ModelOpt 0.44 finds\n", + " # transformer/modelopt_state.pth and restores the FP8 quantizer graph.\n", + " # Modular load_components uses repo+subfolder; ModelOpt ignores subfolder.\n", + " from diffusers import AutoModel\n", + "\n", + " ckpt_dir = materialize_checkpoint_dir(pretrained_source, from_kwargs)\n", + " transformer_dir = ckpt_dir / \"transformer\"\n", + " if not (transformer_dir / \"modelopt_state.pth\").is_file():\n", + " raise RuntimeError(\n", + " f\"FP8 requested but no ModelOpt state at {transformer_dir / 'modelopt_state.pth'}. \"\n", + " \"Point COSMOS3_FP8_MODEL_PATH at a pre-quantized distilled export directory.\"\n", + " )\n", + " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", + " pipe.update_components(transformer=transformer)\n", + " pipe.load_components(\n", + " pretrained_model_name_or_path=distilled_components_root(ckpt_dir),\n", + " torch_dtype=torch.bfloat16,\n", + " )\n", + " else:\n", + " pipe.load_components(torch_dtype=torch.bfloat16)\n", " pipe.enable_safety_checker()\n", " else:\n", " pipe = Cosmos3OmniPipeline.from_pretrained(\n", @@ -918,6 +975,8 @@ "\n", "With FP8 enabled, this notebook loads the public model id from `MODEL_IDS` with `revision=fp8` (override with `COSMOS3_FP8_REVISION`). Optionally set `COSMOS3_FP8_MODEL_PATH` to a local Diffusers checkpoint directory instead.\n", "\n", + "For **4-Step distilled FP8**, ModelOpt 0.44 does not honor Hub `subfolder=` when restoring `modelopt_state.pth`. The helper therefore materializes a local checkpoint root, pre-loads `transformer/` via its full path, pins it with `update_components`, then loads the remaining modular components. After load you can sanity-check restore with `verify_fp8(\"Cosmos3-Super-Text2Image-4Step\")` (or the I2V 4-Step id).\n", + "\n", "Loading may print a `Some weights of the model checkpoint ... not used` warning for `weight_scale` / `input_scale` keys. Those are vLLM-Omni scale tensors shipped alongside the Diffusers ModelOpt state; the warning is expected.\n" ] }, From 619afb87978f0ac1efa2cdcd5d030796cb27c4b8 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Tue, 15 Sep 2026 14:57:53 -0700 Subject: [PATCH 3/7] cookbook(diffusers): point FP8 loads at dedicated Hub repos Load Nano/Super/Edge FP8 from nvidia/Cosmos3-*-FP8 and keep Super 4-Step on revision=fp8. Rebase the notebook onto current audiovisual helpers without documenting mixed precision. --- .../audiovisual/run_with_diffusers.ipynb | 437 +++++++++++------- 1 file changed, 280 insertions(+), 157 deletions(-) diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index d83734e9..f95c69e4 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -17,11 +17,11 @@ "\n", "This notebook runs Cosmos3 audiovisual generation directly with `Cosmos3OmniPipeline`.\n", "\n", - "Run the Cosmos3-Nano, Cosmos3-Super, or Cosmos3-Edge examples independently. Cosmos3-Edge has no audio modules and uses its 480p generation settings. Run Cosmos3-Super T2V/I2V examples without audio. Each section loads the matching model explicitly.\n", + "Run the Cosmos3-Nano, Cosmos3-Super, or Cosmos3-Edge examples independently. Cosmos3-Edge has no audio modules and uses its 480p generation settings. The notebook also includes Nano and Super video-to-video examples, while Super T2V/I2V examples run without audio. Each section loads the matching model explicitly.\n", "\n", - "**FP8:** `Cosmos3-Nano`, `Cosmos3-Super`, and the Super 4-Step distilled checkpoints also ship static-scale FP8 builds. Enable them with `COSMOS3_DIFFUSERS_PRECISION=fp8` (see **Quantized Checkpoints** below). Request cells stay the same; only checkpoint loading changes. `Cosmos3-Edge` has no FP8 build.\n", + "Note: if you have already completed steps 1-3 and installed the `Cosmos3 Diffusers (Python 3.13)` kernel, switch to that kernel and jump directly to step 4. Run the restore cell there, then continue with verification and the examples.\n", "\n", - "Note: if you have already completed steps 1-3 and installed the `Cosmos3 Diffusers (Python 3.13)` kernel, switch to that kernel and jump directly to step 4. Run the restore cell there, then continue with verification and the examples.\n" + "FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano, Super, and Edge use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8`, `nvidia/Cosmos3-Super-FP8`, and `nvidia/Cosmos3-Edge-FP8`. Super 4-Step T2I/I2V keep the existing model cards with `revision=fp8`. Use-case cells stay the same.\n" ] }, { @@ -59,8 +59,7 @@ "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", - "# Optional: bf16 (default) or fp8 for Nano/Super/4-Step — see Quantized Checkpoints\n", - "export COSMOS3_DIFFUSERS_PRECISION=bf16\n", + "export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8\n", "```\n" ] }, @@ -390,44 +389,12 @@ " \"prompt\": \"assets/prompts/text2image/robot_draping.json\",\n", " \"enable_sound\": False,\n", " },\n", - " \"t2i_edge\": {\n", - " \"model\": \"Cosmos3-Edge\",\n", - " \"mode\": \"text2image\",\n", - " \"prompt\": \"assets/prompts/text2image/robot_draping.json\",\n", - " \"enable_sound\": False,\n", - " \"sampling\": {\n", - " \"num_steps\": 50,\n", - " \"guidance\": 5.0,\n", - " \"shift\": 3.0,\n", - " \"fps\": 24,\n", - " \"num_frames\": 1,\n", - " \"resolution\": \"480\",\n", - " \"aspect_ratio\": \"16,9\",\n", - " \"seed\": 0,\n", - " },\n", - " },\n", " \"t2v_nano_noaudio\": {\n", " \"model\": \"Cosmos3-Nano\",\n", " \"mode\": \"text2video\",\n", " \"prompt\": \"assets/prompts/text2video/robot_kitchen.json\",\n", " \"enable_sound\": False,\n", " },\n", - " \"t2v_edge_noaudio\": {\n", - " \"model\": \"Cosmos3-Edge\",\n", - " \"mode\": \"text2video\",\n", - " \"prompt\": \"assets/prompts/text2video/robot_kitchen.json\",\n", - " \"enable_sound\": False,\n", - " \"sampling\": {\n", - " \"num_steps\": 50,\n", - " \"guidance\": 5.0,\n", - " \"shift\": 3.0,\n", - " \"fps\": 24,\n", - " \"num_frames\": 121,\n", - " \"resolution\": \"480\",\n", - " \"aspect_ratio\": \"16,9\",\n", - " \"seed\": 0,\n", - " },\n", - " },\n", " \"t2vs\": {\n", " \"model\": \"Cosmos3-Nano\",\n", " \"mode\": \"text2video\",\n", @@ -444,13 +411,13 @@ " \"i2v_edge_noaudio\": {\n", " \"model\": \"Cosmos3-Edge\",\n", " \"mode\": \"image2video\",\n", - " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving_edge.json\",\n", " \"image\": \"assets/images/image2video/car_driving.jpg\",\n", " \"enable_sound\": False,\n", " \"sampling\": {\n", - " \"num_steps\": 50,\n", + " \"num_steps\": 20,\n", " \"guidance\": 5.0,\n", - " \"shift\": 3.0,\n", + " \"shift\": 8.0,\n", " \"fps\": 24,\n", " \"num_frames\": 121,\n", " \"resolution\": \"480\",\n", @@ -465,6 +432,22 @@ " \"image\": \"assets/images/image2video/coastal_road_audio.jpg\",\n", " \"enable_sound\": True,\n", " },\n", + " \"v2v_nano_noaudio\": {\n", + " \"model\": \"Cosmos3-Nano\",\n", + " \"mode\": \"video2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"negative_prompt\": \"assets/negative_prompts/image2video/neg_prompt.json\",\n", + " \"video\": \"assets/videos/car_driving_plain.mp4\",\n", + " \"enable_sound\": False,\n", + " },\n", + " \"v2vs\": {\n", + " \"model\": \"Cosmos3-Nano\",\n", + " \"mode\": \"video2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"negative_prompt\": \"assets/negative_prompts/image2video/neg_prompt.json\",\n", + " \"video\": \"assets/videos/car_driving_plain.mp4\",\n", + " \"enable_sound\": True,\n", + " },\n", " \"t2v_super_noaudio\": {\n", " \"model\": \"Cosmos3-Super\",\n", " \"mode\": \"text2video\",\n", @@ -478,6 +461,22 @@ " \"image\": \"assets/images/image2video/car_driving.jpg\",\n", " \"enable_sound\": False,\n", " },\n", + " \"v2v_super_noaudio\": {\n", + " \"model\": \"Cosmos3-Super\",\n", + " \"mode\": \"video2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"negative_prompt\": \"assets/negative_prompts/image2video/neg_prompt.json\",\n", + " \"video\": \"assets/videos/car_driving_plain.mp4\",\n", + " \"enable_sound\": False,\n", + " },\n", + " \"v2vs_super\": {\n", + " \"model\": \"Cosmos3-Super\",\n", + " \"mode\": \"video2video\",\n", + " \"prompt\": \"assets/prompts/image2video/car_driving.json\",\n", + " \"negative_prompt\": \"assets/negative_prompts/image2video/neg_prompt.json\",\n", + " \"video\": \"assets/videos/car_driving_plain.mp4\",\n", + " \"enable_sound\": True,\n", + " },\n", " # Distilled 4-step checkpoints bake classifier-free guidance into the weights and read their\n", " # sigma schedule from the checkpoint, so they carry no num_steps/guidance/shift.\n", " \"t2i_distilled\": {\n", @@ -554,7 +553,10 @@ " prompt_path = asset_path(spec[\"prompt\"])\n", " negative_prompt = \"\"\n", " if spec[\"mode\"] != \"text2image\":\n", - " negative_prompt_path = asset_path(f\"assets/negative_prompts/{spec['mode']}/neg_prompt.json\")\n", + " # There is no negative_prompts/video2video directory; those cases name the file they reuse.\n", + " negative_prompt_path = asset_path(\n", + " spec.get(\"negative_prompt\") or f\"assets/negative_prompts/{spec['mode']}/neg_prompt.json\"\n", + " )\n", " negative_prompt = compact_json_file(negative_prompt_path)\n", " payload_path = payload_dir / f\"{use_case}.json\"\n", " payload = {\n", @@ -577,6 +579,11 @@ " if spec[\"mode\"] == \"image2video\":\n", " image_path = asset_path(spec[\"image\"])\n", " payload[\"vision_path\"] = os.path.relpath(image_path, payload_path.parent)\n", + " if spec[\"mode\"] == \"video2video\":\n", + " video_path = asset_path(spec[\"video\"])\n", + " payload[\"vision_path\"] = os.path.relpath(video_path, payload_path.parent)\n", + " payload[\"condition_frame_indexes_vision\"] = spec.get(\"condition_frame_indexes_vision\", [0, 1])\n", + " payload[\"condition_video_keep\"] = spec.get(\"condition_video_keep\", \"first\")\n", "\n", " payload_path.write_text(json.dumps(payload, indent=2) + \"\\n\")\n", "\n", @@ -588,9 +595,14 @@ " print(f\"output: {output_dir}\")\n", " print(f\"prompt: {prompt_path.relative_to(COSMOS_ROOT)}\")\n", " if \"vision_path\" in payload:\n", - " image_display_path = resolve_payload_path(payload_path, payload[\"vision_path\"])\n", - " print(f\"image: {image_display_path.relative_to(COSMOS_ROOT)}\")\n", - " display(Image(filename=str(image_display_path), width=420))\n", + " vision_display_path = resolve_payload_path(payload_path, payload[\"vision_path\"])\n", + " if payload[\"model_mode\"] == \"video2video\":\n", + " print(f\"video: {vision_display_path.relative_to(COSMOS_ROOT)}\")\n", + " display_video(vision_display_path, width=420)\n", + " else:\n", + " print(f\"image: {vision_display_path.relative_to(COSMOS_ROOT)}\")\n", + " display(Image(filename=str(vision_display_path), width=420))\n", + " print(json.dumps({k: payload[k] for k in [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]}, indent=2))\n", " summary_keys = [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]\n", " print(json.dumps({k: payload[k] for k in summary_keys if k in payload}, indent=2))\n", " if payload.get(\"distilled\"):\n", @@ -616,7 +628,7 @@ "from diffusers import Cosmos3OmniPipeline\n", "from diffusers import logging as diffusers_logging\n", "from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler\n", - "from diffusers.utils import encode_video, export_to_video, load_image\n", + "from diffusers.utils import encode_video, export_to_video, load_image, load_video\n", "\n", "try:\n", " from diffusers import Cosmos3DistilledModularPipeline\n", @@ -638,14 +650,16 @@ " \"Cosmos3-Super-Image2Video-4Step\",\n", "}\n", "\n", - "# FP8 builds for Nano/Super and Super 4-Step (same request path as bf16; see Quantized Checkpoints).\n", - "FP8_SUPPORTED_MODELS = {\n", - " \"Cosmos3-Nano\",\n", - " \"Cosmos3-Super\",\n", - " \"Cosmos3-Super-Text2Image-4Step\",\n", - " \"Cosmos3-Super-Image2Video-4Step\",\n", + "# Base models ship FP8 on dedicated Hub repos. Distilled 4-Step FP8 stays on revision=fp8\n", + "# of the existing cards.\n", + "FP8_SEPARATE_REPOS = {\n", + " \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n", + " \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n", + " \"Cosmos3-Edge\": \"nvidia/Cosmos3-Edge-FP8\",\n", "}\n", "\n", + "FP8_SUPPORTED_MODELS = set(FP8_SEPARATE_REPOS) | set(DISTILLED_MODELS)\n", + "\n", "_pipe = None\n", "_pipe_model = None\n", "\n", @@ -658,20 +672,29 @@ " return precision_mode() == \"fp8\"\n", "\n", "\n", - "def resolve_fp8_pretrained_source(model: str):\n", - " \"\"\"Return a local path or public HF repo id for the FP8 checkpoint.\n", + "def resolve_model_id(model: str) -> str:\n", + " return MODEL_IDS.get(model, model)\n", + "\n", "\n", - " Uses `MODEL_IDS[model]` with `revision=COSMOS3_FP8_REVISION` (default `fp8`) unless\n", - " `COSMOS3_FP8_MODEL_PATH` points at a local Diffusers checkpoint directory.\n", - " \"\"\"\n", + "def resolve_pretrained_source(model: str):\n", + " \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n", " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", - " if override:\n", + " if is_fp8_precision() and override:\n", " return Path(override).expanduser().resolve()\n", + " if is_fp8_precision() and model in FP8_SEPARATE_REPOS:\n", + " return FP8_SEPARATE_REPOS[model]\n", " return resolve_model_id(model)\n", "\n", "\n", - "def resolve_model_id(model: str) -> str:\n", - " return MODEL_IDS.get(model, model)\n", + "def pretrained_from_kwargs(model: str) -> dict:\n", + " kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", + " if (\n", + " is_fp8_precision()\n", + " and not os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " and model in DISTILLED_MODELS\n", + " ):\n", + " kwargs[\"revision\"] = os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\")\n", + " return kwargs\n", "\n", "\n", "def is_distilled_model(model: str) -> bool:\n", @@ -751,9 +774,9 @@ "\n", "def get_pipe(model: str):\n", " global _pipe, _pipe_model\n", - " model_id = resolve_model_id(model)\n", " prec = precision_mode()\n", - " cache_key = f\"{model_id}@{prec}\"\n", + " pretrained_source = resolve_pretrained_source(model)\n", + " cache_key = f\"{pretrained_source}@{prec}\"\n", " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", @@ -765,21 +788,16 @@ " f\"FP8 is only supported for {sorted(FP8_SUPPORTED_MODELS)}; got {model}. \"\n", " \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n", " )\n", + " from_kwargs = pretrained_from_kwargs(model)\n", " if fp8:\n", " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", " from modelopt.torch.opt import enable_huggingface_checkpointing\n", "\n", " enable_huggingface_checkpointing()\n", - " pretrained_source = resolve_fp8_pretrained_source(model)\n", - " from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", - " if not os.environ.get(\"COSMOS3_FP8_MODEL_PATH\"):\n", - " from_kwargs[\"revision\"] = os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\")\n", " print(f\"loading FP8 {pretrained_source}...\")\n", " else:\n", - " pretrained_source = model_id\n", - " from_kwargs = {\"token\": os.environ.get(\"HF_TOKEN\") or None}\n", - " print(f\"loading {model_id}...\")\n", + " print(f\"loading {pretrained_source}...\")\n", "\n", " if is_distilled_model(model):\n", " if Cosmos3DistilledModularPipeline is None:\n", @@ -840,7 +858,9 @@ " if not distilled:\n", " pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=payload[\"shift\"])\n", " generator = torch.Generator(device=\"cuda\").manual_seed(payload[\"seed\"])\n", - " image = load_image(str(resolve_payload_path(payload_path, payload[\"vision_path\"]))) if payload[\"model_mode\"] == \"image2video\" else None\n", + " vision_path = resolve_payload_path(payload_path, payload[\"vision_path\"]) if \"vision_path\" in payload else None\n", + " image = load_image(str(vision_path)) if payload[\"model_mode\"] == \"image2video\" else None\n", + " video = load_video(str(vision_path)) if payload[\"model_mode\"] == \"video2video\" else None\n", " output_ext = \".png\" if payload[\"model_mode\"] == \"text2image\" else \".mp4\"\n", " output_path = output_dir / f\"{payload['name']}{output_ext}\"\n", "\n", @@ -893,6 +913,9 @@ " prompt=payload[\"prompt\"],\n", " negative_prompt=payload[\"negative_prompt\"],\n", " image=image,\n", + " video=video,\n", + " condition_frame_indexes_vision=payload.get(\"condition_frame_indexes_vision\", [0, 1]),\n", + " condition_video_keep=payload.get(\"condition_video_keep\", \"first\"),\n", " num_frames=payload[\"num_frames\"],\n", " height=height,\n", " width=width,\n", @@ -958,42 +981,31 @@ }, { "cell_type": "markdown", - "id": "quantized-checkpoints-fp8", + "id": "fp8-quantized-checkpoints", "metadata": {}, "source": [ "## Quantized Checkpoints (FP8)\n", "\n", - "FP8 builds are available for **`Cosmos3-Nano`**, **`Cosmos3-Super`**, and the Super **4-Step** distilled checkpoints (`Cosmos3-Edge` has no FP8 build). Serving FP8 is a checkpoint-loading change: the use-case request cells below stay the same.\n", - "\n", - "Static-scale FP8 checkpoints are produced with [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer) and load through the same Diffusers pipelines as bf16 (`Cosmos3OmniPipeline` for Nano/Super, `Cosmos3DistilledModularPipeline` for 4-Step). The install cell pins `nvidia-modelopt==0.44.0`.\n", + "Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before running the helper cell's `get_pipe` (or re-run that cell after changing it). Use-case request cells stay the same.\n", "\n", - "Enable FP8 before running the helper cell (or re-run the helper cell after changing the env):\n", - "\n", - "```bash\n", - "export COSMOS3_DIFFUSERS_PRECISION=fp8\n", - "```\n", + "- **Nano / Super / Edge:** load `nvidia/Cosmos3-Nano-FP8`, `nvidia/Cosmos3-Super-FP8`, and `nvidia/Cosmos3-Edge-FP8` (default revision on those repos). Do not pass `revision=fp8` on the BF16 cards.\n", + "- **Super 4-Step T2I / I2V:** load the existing cards with `revision=fp8`.\n", "\n", - "With FP8 enabled, this notebook loads the public model id from `MODEL_IDS` with `revision=fp8` (override with `COSMOS3_FP8_REVISION`). Optionally set `COSMOS3_FP8_MODEL_PATH` to a local Diffusers checkpoint directory instead.\n", + "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory. Distilled FP8 preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`. After load, `verify_fp8(\"Cosmos3-Nano\")` (or another supported model) checks FP8 weights and quantizer modules.\n", "\n", - "For **4-Step distilled FP8**, ModelOpt 0.44 does not honor Hub `subfolder=` when restoring `modelopt_state.pth`. The helper therefore materializes a local checkpoint root, pre-loads `transformer/` via its full path, pins it with `update_components`, then loads the remaining modular components. After load you can sanity-check restore with `verify_fp8(\"Cosmos3-Super-Text2Image-4Step\")` (or the I2V 4-Step id).\n", - "\n", - "Loading may print a `Some weights of the model checkpoint ... not used` warning for `weight_scale` / `input_scale` keys. Those are vLLM-Omni scale tensors shipped alongside the Diffusers ModelOpt state; the warning is expected.\n" + "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" ] }, { "cell_type": "code", - "execution_count": null, - "id": "quantized-checkpoints-fp8-enable", + "id": "fp8-precision-toggle", "metadata": {}, "outputs": [], + "execution_count": null, "source": [ - "# Optional: flip precision for subsequent get_pipe / run_diffusers_payload calls.\n", - "# Re-run the helper cell above after changing this if helpers were already executed.\n", - "import os\n", - "\n", - "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\" # uncomment to use FP8 Nano/Super/4-Step\n", - "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n", - "print(\"COSMOS3_FP8_REVISION:\", os.environ.get(\"COSMOS3_FP8_REVISION\", \"fp8\"))\n" + "# Default is bf16. Uncomment to load FP8 checkpoints in the cells below.\n", + "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\"\n", + "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n" ] }, { @@ -1356,7 +1368,122 @@ }, { "cell_type": "markdown", - "id": "9f58eb8b", + "id": "v2v-nano-noaudio-md", + "metadata": {}, + "source": [ + "## Nano: Video to Video Without Audio\n", + "\n", + "Nano video-to-video generation from a reference clip, with audio disabled. The first two latent frames of the reference clip are kept as clean conditioning and the rest are denoised.\n", + "\n", + "### Create Payload\n" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "v2v_nano_noaudio_payload, v2v_nano_noaudio_output, v2v_nano_noaudio_model = create_payload(\"v2v_nano_noaudio\", backend=\"diffusers\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-nano-noaudio-payload" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run\n" + ], + "id": "v2v-nano-noaudio-run-md" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run_diffusers_payload(v2v_nano_noaudio_payload, v2v_nano_noaudio_output, model=\"Cosmos3-Nano\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-nano-noaudio-run" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### View Results\n" + ], + "id": "v2v-nano-noaudio-view-md" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "view_run(v2v_nano_noaudio_output)\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-nano-noaudio-view" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Nano: Video to Video with Audio\n", + "\n", + "Nano video-to-video generation from a reference clip, with generated audio. The first two latent frames of the reference clip are kept as clean conditioning and the rest are denoised.\n", + "\n", + "### Create Payload\n" + ], + "id": "v2vs-md" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "v2vs_payload, v2vs_output, v2vs_model = create_payload(\"v2vs\", backend=\"diffusers\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-payload" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run\n" + ], + "id": "v2vs-run-md" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "run_diffusers_payload(v2vs_payload, v2vs_output, model=\"Cosmos3-Nano\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-run" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### View Results\n" + ], + "id": "v2vs-view-md" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "view_run(v2vs_output)\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-view" + }, + { + "cell_type": "markdown", "metadata": {}, "source": [ "## Super: Text to Video Without Audio\n", @@ -1570,20 +1697,16 @@ "cell_type": "markdown", "id": "482d68de", "metadata": {}, - "source": [ - "# Cosmos3-Edge Examples\n", - "\n", - "Cosmos3-Edge runs on a single GPU and has no audio tokenizer, so all Edge video examples keep audio disabled. The Edge checkpoint is trained for 480p generation; these examples use 832x480, 121 frames, 50 denoising steps, guidance scale 5.0, and flow shift 3.0.\n" - ] + "source": "# Cosmos3-Edge Example\n\nCosmos3-Edge runs on a single GPU and has no audio tokenizer, so the Edge video example keeps audio disabled. The Edge checkpoint is trained for 480p generation; this example uses 832x480, 121 frames, 20 denoising steps, guidance scale 5.0, and flow shift 8.0.\n" }, { "cell_type": "markdown", - "id": "c5e29c2a", + "id": "33d5e6ba", "metadata": {}, "source": [ - "## Edge: Text to Image\n", + "## Edge: Image to Video Without Audio\n", "\n", - "Generate a single 832x480 image from the structured JSON prompt.\n", + "Animate the supplied driving image for 121 frames at 832x480.\n", "\n", "### Create Payload\n" ] @@ -1591,16 +1714,16 @@ { "cell_type": "code", "execution_count": null, - "id": "edd2c2ed", + "id": "b4dcd641", "metadata": {}, "outputs": [], "source": [ - "t2i_edge_payload, t2i_edge_output, t2i_edge_model = create_payload(\"t2i_edge\", backend=\"diffusers\")\n" + "i2v_edge_noaudio_payload, i2v_edge_noaudio_output, i2v_edge_noaudio_model = create_payload(\"i2v_edge_noaudio\", backend=\"diffusers\")\n" ] }, { "cell_type": "markdown", - "id": "444115f2", + "id": "2043536b", "metadata": {}, "source": [ "### Run\n" @@ -1609,16 +1732,16 @@ { "cell_type": "code", "execution_count": null, - "id": "65857776", + "id": "aa8a12d3", "metadata": {}, "outputs": [], "source": [ - "run_diffusers_payload(t2i_edge_payload, t2i_edge_output, model=\"Cosmos3-Edge\")\n" + "run_diffusers_payload(i2v_edge_noaudio_payload, i2v_edge_noaudio_output, model=\"Cosmos3-Edge\")\n" ] }, { "cell_type": "markdown", - "id": "9c1d7b62", + "id": "a022c344", "metadata": {}, "source": [ "### View Results\n" @@ -1627,128 +1750,128 @@ { "cell_type": "code", "execution_count": null, - "id": "d4cd4817", + "id": "c40acf67", "metadata": {}, "outputs": [], "source": [ - "view_run(t2i_edge_output)\n" + "view_run(i2v_edge_noaudio_output)\n" ] }, { "cell_type": "markdown", - "id": "2fb4e265", "metadata": {}, "source": [ - "## Edge: Text to Video Without Audio\n", + "## Super: Video to Video Without Audio\n", "\n", - "Generate a 121-frame 832x480 video from the structured JSON prompt.\n", + "Super video-to-video generation using the same reference clip and conditioning frames as Nano, with audio disabled.\n", "\n", "### Create Payload\n" - ] + ], + "id": "v2v-super-md" }, { "cell_type": "code", - "execution_count": null, - "id": "c133f313", "metadata": {}, - "outputs": [], "source": [ - "t2v_edge_noaudio_payload, t2v_edge_noaudio_output, t2v_edge_noaudio_model = create_payload(\"t2v_edge_noaudio\", backend=\"diffusers\")\n" - ] + "v2v_super_noaudio_payload, v2v_super_noaudio_output, v2v_super_noaudio_model = create_payload(\"v2v_super_noaudio\", backend=\"diffusers\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-super-payload" }, { "cell_type": "markdown", - "id": "490c2c93", "metadata": {}, "source": [ "### Run\n" - ] + ], + "id": "v2v-super-run-md" }, { "cell_type": "code", - "execution_count": null, - "id": "557831a9", "metadata": {}, - "outputs": [], "source": [ - "run_diffusers_payload(t2v_edge_noaudio_payload, t2v_edge_noaudio_output, model=\"Cosmos3-Edge\")\n" - ] + "run_diffusers_payload(v2v_super_noaudio_payload, v2v_super_noaudio_output, model=\"Cosmos3-Super\")\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-super-run" }, { "cell_type": "markdown", - "id": "6fdd64ed", "metadata": {}, "source": [ "### View Results\n" - ] + ], + "id": "v2v-super-view-md" }, { "cell_type": "code", - "execution_count": null, - "id": "44f6f2e8", "metadata": {}, - "outputs": [], "source": [ - "view_run(t2v_edge_noaudio_output)\n" - ] + "view_run(v2v_super_noaudio_output)\n" + ], + "execution_count": null, + "outputs": [], + "id": "v2v-super-view" }, { "cell_type": "markdown", - "id": "33d5e6ba", "metadata": {}, "source": [ - "## Edge: Image to Video Without Audio\n", + "## Super: Video to Video with Audio\n", "\n", - "Animate the supplied driving image for 121 frames at 832x480.\n", + "Super video-to-video generation on the same reference clip and conditioning frames as the Nano case, with generated audio.\n", "\n", - "### Create Payload\n" - ] + "### Create Payload" + ], + "id": "v2vs-super-md" }, { "cell_type": "code", - "execution_count": null, - "id": "b4dcd641", "metadata": {}, - "outputs": [], "source": [ - "i2v_edge_noaudio_payload, i2v_edge_noaudio_output, i2v_edge_noaudio_model = create_payload(\"i2v_edge_noaudio\", backend=\"diffusers\")\n" - ] + "v2vs_super_payload, v2vs_super_output, v2vs_super_model = create_payload(\"v2vs_super\", backend=\"diffusers\")" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-super-payload" }, { "cell_type": "markdown", - "id": "2043536b", "metadata": {}, "source": [ - "### Run\n" - ] + "### Run" + ], + "id": "v2vs-super-run-md" }, { "cell_type": "code", - "execution_count": null, - "id": "aa8a12d3", "metadata": {}, - "outputs": [], "source": [ - "run_diffusers_payload(i2v_edge_noaudio_payload, i2v_edge_noaudio_output, model=\"Cosmos3-Edge\")\n" - ] + "run_diffusers_payload(v2vs_super_payload, v2vs_super_output, model=\"Cosmos3-Super\")" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-super-run" }, { "cell_type": "markdown", - "id": "a022c344", "metadata": {}, "source": [ - "### View Results\n" - ] + "### View Results" + ], + "id": "v2vs-super-view-md" }, { "cell_type": "code", - "execution_count": null, - "id": "c40acf67", "metadata": {}, - "outputs": [], "source": [ - "view_run(i2v_edge_noaudio_output)\n" - ] + "view_run(v2vs_super_output)" + ], + "execution_count": null, + "outputs": [], + "id": "v2vs-super-view" } ], "metadata": { From c4ecc2dbb6b94b37be44fcb12d1e2bac22dd247b Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Thu, 17 Sep 2026 12:30:03 -0700 Subject: [PATCH 4/7] cookbook(diffusers): mention checkpoint mixed W8A8/W8A16 schedule Point users at the Cosmos3 Diffusers docs instead of duplicating generate-time kwargs. --- .../cosmos3/generator/audiovisual/run_with_diffusers.ipynb | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index f95c69e4..3e81f566 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -991,6 +991,8 @@ "- **Nano / Super / Edge:** load `nvidia/Cosmos3-Nano-FP8`, `nvidia/Cosmos3-Super-FP8`, and `nvidia/Cosmos3-Edge-FP8` (default revision on those repos). Do not pass `revision=fp8` on the BF16 cards.\n", "- **Super 4-Step T2I / I2V:** load the existing cards with `revision=fp8`.\n", "\n", + "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). Distilled 4-Step stays all W8A8. See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", + "\n", "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory. Distilled FP8 preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`. After load, `verify_fp8(\"Cosmos3-Nano\")` (or another supported model) checks FP8 weights and quantizer modules.\n", "\n", "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" From b8fdd60bc171e2e750b1c93c34c0619dc0b2ff86 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 28 Sep 2026 16:02:56 -0700 Subject: [PATCH 5/7] cookbook(diffusers): load action and transfer FP8 from dedicated Hub repos Point forward dynamics, inverse dynamics, and video transfer at the Nano and Super FP8 model cards without changing their request cells. --- .../action/run_fd_with_diffusers.ipynb | 102 +++++++++++-- .../action/run_id_with_diffusers.ipynb | 104 +++++++++++-- .../run_video_transfer_with_diffusers.ipynb | 140 ++++++++++++++++-- 3 files changed, 307 insertions(+), 39 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb index 2ae342f8..ebc76b2c 100644 --- a/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb @@ -20,7 +20,9 @@ "\n", "It runs three AV ego trajectories, three camera-pose trajectories, an autoregressive multiview DROID rollout, an autoregressive UMI rollout, and a bimanual hand-pose chunk. For inverse dynamics, see [`run_id_with_diffusers.ipynb`](./run_id_with_diffusers.ipynb).\n", "\n", - "> **Guardrails:** Disabled by default in this notebook. Set `COSMOS3_DIFFUSERS_GUARDRAILS=true` before running the helper cell to enable them." + "> **Guardrails:** Disabled by default in this notebook. Set `COSMOS3_DIFFUSERS_GUARDRAILS=true` before running the helper cell to enable them.\n", + "\n", + "FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano and Super use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Use-case cells stay the same.\n" ] }, { @@ -63,7 +65,8 @@ "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", - "```" + "export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8\n", + "```\n" ] }, { @@ -114,6 +117,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -126,6 +130,7 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -133,7 +138,7 @@ " print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -185,7 +190,8 @@ " scipy \\\n", " torch \\\n", " torchvision \\\n", - " transformers\n", + " transformers \\\n", + " \"nvidia-modelopt==0.44.0\"\n", "\n", "# Video decoding uses the PyAV fallback defined in the helpers cell, avoiding TorchCodec ABI\n", "# coupling with the selected PyTorch backend.\n", @@ -197,7 +203,7 @@ "echo\n", "echo \"Installed dependencies into: $COSMOS3_DIFFUSERS_ACTION_VENV\"\n", "echo \"Next: switch this notebook kernel to: Cosmos3 Diffusers Action (Python 3.13)\"\n", - "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"" + "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"\n" ] }, { @@ -262,6 +268,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -274,6 +281,7 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -281,7 +289,7 @@ " print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -571,25 +579,65 @@ " print(f\"released previous pipeline; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", "\n", "\n", + "\n", + "FP8_SEPARATE_REPOS = {\n", + " \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n", + " \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n", + "}\n", + "\n", + "\n", + "def precision_mode() -> str:\n", + " return os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\").strip().lower()\n", + "\n", + "\n", + "def is_fp8_precision() -> bool:\n", + " return precision_mode() == \"fp8\"\n", + "\n", + "\n", + "def resolve_pretrained_source(model: str):\n", + " \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n", + " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " if is_fp8_precision() and override:\n", + " return Path(override).expanduser().resolve()\n", + " if is_fp8_precision():\n", + " if model not in FP8_SEPARATE_REPOS:\n", + " raise RuntimeError(\n", + " f\"FP8 is only supported for {sorted(FP8_SEPARATE_REPOS)}; got {model}. \"\n", + " \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n", + " )\n", + " return FP8_SEPARATE_REPOS[model]\n", + " return MODEL_IDS.get(model, model)\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniPipeline:\n", " global _pipe, _pipe_model\n", - " model_id = MODEL_IDS.get(model, model)\n", - " if _pipe is not None and _pipe_model == model_id:\n", + " prec = precision_mode()\n", + " pretrained_source = resolve_pretrained_source(model)\n", + " cache_key = f\"{pretrained_source}@{prec}\"\n", + " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", - " print(f\"loading {model_id}...\")\n", " t0 = time.time()\n", + " if is_fp8_precision():\n", + " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", + " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", + " from modelopt.torch.opt import enable_huggingface_checkpointing\n", + "\n", + " enable_huggingface_checkpointing()\n", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", " pipe = Cosmos3OmniPipeline.from_pretrained(\n", - " model_id,\n", + " pretrained_source,\n", " torch_dtype=torch.bfloat16,\n", " safety_checker=None,\n", " enable_safety_checker=GUARDRAILS,\n", " token=os.environ.get(\"HF_TOKEN\") or None,\n", " )\n", " pipe.to(\"cuda\")\n", - " _pipe, _pipe_model = pipe, model_id\n", - " print(f\"loaded pipeline in {time.time() - t0:.1f}s; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", + " _pipe, _pipe_model = pipe, cache_key\n", + " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", " return _pipe\n", "\n", "\n", @@ -748,6 +796,36 @@ " return output_path\n" ] }, + { + "cell_type": "markdown", + "id": "fp8-quantized-checkpoints", + "metadata": {}, + "source": [ + "## Quantized Checkpoints (FP8)\n", + "\n", + "Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before running the helper cell's `get_pipe` (or re-run that cell after changing it). Use-case request cells stay the same.\n", + "\n", + "Nano and Super load `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Do not pass `revision=fp8` on the BF16 cards.\n", + "\n", + "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", + "\n", + "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory.\n", + "\n", + "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + ] + }, + { + "cell_type": "code", + "id": "fp8-precision-toggle", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "# Default is bf16. Uncomment to load FP8 checkpoints in the cells below.\n", + "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\"\n", + "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n" + ] + }, { "cell_type": "markdown", "id": "action-diffusers-helpers-15", diff --git a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb index d8417396..647c91fb 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb @@ -18,7 +18,9 @@ "\n", "Inverse dynamics runs the other direction from forward dynamics: the input is a clip and no trajectory, and the model recovers the motion that connects the frames. It runs `Cosmos3OmniPipeline` with the clip grouped into a `CosmosActionCondition`.\n", "\n", - "It runs the two AV clips plus the Bridge, AgiBotWorld-Beta, RoboMIND Franka, RoboMIND Franka dual-arm, RoboMIND UR, UMI, and Fractal episodes under [`assets/`](./assets). For forward dynamics, see [`run_fd_with_diffusers.ipynb`](./run_fd_with_diffusers.ipynb)." + "It runs the two AV clips plus the Bridge, AgiBotWorld-Beta, RoboMIND Franka, RoboMIND Franka dual-arm, RoboMIND UR, UMI, and Fractal episodes under [`assets/`](./assets). For forward dynamics, see [`run_fd_with_diffusers.ipynb`](./run_fd_with_diffusers.ipynb).\n", + "\n", + "FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano and Super use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Use-case cells stay the same.\n" ] }, { @@ -61,7 +63,8 @@ "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", - "```" + "export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8\n", + "```\n" ] }, { @@ -112,6 +115,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -124,6 +128,7 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -131,7 +136,7 @@ " print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -183,7 +188,8 @@ " scipy \\\n", " torch \\\n", " torchvision \\\n", - " transformers\n", + " transformers \\\n", + " \"nvidia-modelopt==0.44.0\"\n", "\n", "# Video decoding uses the PyAV fallback defined in the helpers cell, avoiding TorchCodec ABI\n", "# coupling with the selected PyTorch backend.\n", @@ -195,7 +201,7 @@ "echo\n", "echo \"Installed dependencies into: $COSMOS3_DIFFUSERS_ACTION_VENV\"\n", "echo \"Next: switch this notebook kernel to: Cosmos3 Diffusers Action (Python 3.13)\"\n", - "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"" + "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"\n" ] }, { @@ -260,6 +266,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -272,6 +279,7 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", @@ -279,7 +287,7 @@ " print(f\"note: no cosmos_framework package under {COSMOS3_REPO}; point COSMOS3_REPO at your framework checkout\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -374,7 +382,7 @@ "\n", "Inverse dynamics reads an existing clip and predicts the action chunk that connects its frames, so it conditions on `action.video` and passes no `raw_actions`. The only output is the predicted action array; nothing is generated to view as a video.\n", "\n", - "Predicted AV actions are 9D framewise deltas \u2014 a 3D translation plus a 6D rotation \u2014 which integrate into absolute camera poses for plotting.\n", + "Predicted AV actions are 9D framewise deltas — a 3D translation plus a 6D rotation — which integrate into absolute camera poses for plotting.\n", "\n", "Robot episodes work the same way with two differences: their actions come back normalized, so they are denormalized with the episode's own statistics first, and a single action vector can carry several end-effectors, so each case names the 9D block to integrate. Longer episodes are covered by predicting one chunk per sampled window and concatenating the results." ] @@ -676,25 +684,65 @@ " print(f\"released previous pipeline; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", "\n", "\n", + "\n", + "FP8_SEPARATE_REPOS = {\n", + " \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n", + " \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n", + "}\n", + "\n", + "\n", + "def precision_mode() -> str:\n", + " return os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\").strip().lower()\n", + "\n", + "\n", + "def is_fp8_precision() -> bool:\n", + " return precision_mode() == \"fp8\"\n", + "\n", + "\n", + "def resolve_pretrained_source(model: str):\n", + " \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n", + " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " if is_fp8_precision() and override:\n", + " return Path(override).expanduser().resolve()\n", + " if is_fp8_precision():\n", + " if model not in FP8_SEPARATE_REPOS:\n", + " raise RuntimeError(\n", + " f\"FP8 is only supported for {sorted(FP8_SEPARATE_REPOS)}; got {model}. \"\n", + " \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n", + " )\n", + " return FP8_SEPARATE_REPOS[model]\n", + " return MODEL_IDS.get(model, model)\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniPipeline:\n", " global _pipe, _pipe_model\n", - " model_id = MODEL_IDS.get(model, model)\n", - " if _pipe is not None and _pipe_model == model_id:\n", + " prec = precision_mode()\n", + " pretrained_source = resolve_pretrained_source(model)\n", + " cache_key = f\"{pretrained_source}@{prec}\"\n", + " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", - " print(f\"loading {model_id}...\")\n", " t0 = time.time()\n", + " if is_fp8_precision():\n", + " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", + " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", + " from modelopt.torch.opt import enable_huggingface_checkpointing\n", + "\n", + " enable_huggingface_checkpointing()\n", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", " pipe = Cosmos3OmniPipeline.from_pretrained(\n", - " model_id,\n", + " pretrained_source,\n", " torch_dtype=torch.bfloat16,\n", " safety_checker=None,\n", " enable_safety_checker=GUARDRAILS,\n", " token=os.environ.get(\"HF_TOKEN\") or None,\n", " )\n", " pipe.to(\"cuda\")\n", - " _pipe, _pipe_model = pipe, model_id\n", - " print(f\"loaded pipeline in {time.time() - t0:.1f}s; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", + " _pipe, _pipe_model = pipe, cache_key\n", + " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", " return _pipe\n", "\n", "\n", @@ -915,6 +963,36 @@ " )\n" ] }, + { + "cell_type": "markdown", + "id": "fp8-quantized-checkpoints", + "metadata": {}, + "source": [ + "## Quantized Checkpoints (FP8)\n", + "\n", + "Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before running the helper cell's `get_pipe` (or re-run that cell after changing it). Use-case request cells stay the same.\n", + "\n", + "Nano and Super load `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Do not pass `revision=fp8` on the BF16 cards.\n", + "\n", + "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", + "\n", + "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory.\n", + "\n", + "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + ] + }, + { + "cell_type": "code", + "id": "fp8-precision-toggle", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "# Default is bf16. Uncomment to load FP8 checkpoints in the cells below.\n", + "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\"\n", + "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n" + ] + }, { "cell_type": "markdown", "id": "action-diffusers-usecases", diff --git a/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb index 081a3db6..027bb764 100644 --- a/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb @@ -18,7 +18,9 @@ "\n", "This notebook runs Cosmos3 structural-control transfer with `Cosmos3OmniModularPipeline`, the pipeline that exposes control inputs.\n", "\n", - "It runs edge, blur, depth, segmentation, and world-scenario-map transfer on Cosmos3-Nano, then the same five controls on Cosmos3-Super." + "It runs edge, blur, depth, segmentation, and world-scenario-map transfer on Cosmos3-Nano, then the same five controls on Cosmos3-Super.\n", + "\n", + "FP8 checkpoints: set `COSMOS3_DIFFUSERS_PRECISION=fp8` before loading. Nano and Super use the dedicated Hub repos `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Use-case cells stay the same.\n" ] }, { @@ -58,7 +60,8 @@ "export HF_HOME=/path/to/large/huggingface/cache\n", "export UV_LINK_MODE=copy\n", "export CUDA_VISIBLE_DEVICES=0\n", - "```" + "export COSMOS3_DIFFUSERS_PRECISION=bf16 # or fp8\n", + "```\n" ] }, { @@ -108,6 +111,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -120,12 +124,13 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -165,7 +170,8 @@ " ipykernel \\\n", " torch \\\n", " torchvision \\\n", - " transformers\n", + " transformers \\\n", + " \"nvidia-modelopt==0.44.0\"\n", "\n", "\"$COSMOS3_DIFFUSERS_VENV/bin/python\" -m ipykernel install --user \\\n", " --name cosmos3-diffusers \\\n", @@ -174,7 +180,7 @@ "echo\n", "echo \"Installed dependencies into: $COSMOS3_DIFFUSERS_VENV\"\n", "echo \"Next: switch this notebook kernel to: Cosmos3 Diffusers (Python 3.13)\"\n", - "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"" + "echo \"After switching kernels, run the Restore Environment cell below, then continue with Verify.\"\n" ] }, { @@ -238,6 +244,7 @@ " os.environ.setdefault(\"HF_HOME\", str(Path.home() / \".cache\" / \"huggingface\"))\n", " os.environ.setdefault(\"HF_HUB_DISABLE_XET\", \"1\")\n", " os.environ.setdefault(\"CUDA_VISIBLE_DEVICES\", \"0\")\n", + " os.environ.setdefault(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\")\n", "\n", " print(f\"COSMOS_ROOT: {COSMOS_ROOT}\")\n", " for key in [\n", @@ -250,12 +257,13 @@ " \"HF_HOME\",\n", " \"HF_HUB_DISABLE_XET\",\n", " \"CUDA_VISIBLE_DEVICES\",\n", + " \"COSMOS3_DIFFUSERS_PRECISION\",\n", " ]:\n", " print(f\"{key}: {os.environ[key]}\")\n", " print(\"HF_TOKEN:\", \"\" if os.environ.get(\"HF_TOKEN\") else \"\")\n", "\n", "\n", - "configure_diffusers_environment()" + "configure_diffusers_environment()\n" ] }, { @@ -435,19 +443,93 @@ " print(f\"released previous pipeline; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", "\n", "\n", + "\n", + "FP8_SEPARATE_REPOS = {\n", + " \"Cosmos3-Nano\": \"nvidia/Cosmos3-Nano-FP8\",\n", + " \"Cosmos3-Super\": \"nvidia/Cosmos3-Super-FP8\",\n", + "}\n", + "\n", + "\n", + "def precision_mode() -> str:\n", + " return os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\").strip().lower()\n", + "\n", + "\n", + "def is_fp8_precision() -> bool:\n", + " return precision_mode() == \"fp8\"\n", + "\n", + "\n", + "def resolve_pretrained_source(model: str):\n", + " \"\"\"Hub repo id or local path for the requested precision.\"\"\"\n", + " override = os.environ.get(\"COSMOS3_FP8_MODEL_PATH\")\n", + " if is_fp8_precision() and override:\n", + " return Path(override).expanduser().resolve()\n", + " if is_fp8_precision():\n", + " if model not in FP8_SEPARATE_REPOS:\n", + " raise RuntimeError(\n", + " f\"FP8 is only supported for {sorted(FP8_SEPARATE_REPOS)}; got {model}. \"\n", + " \"Unset COSMOS3_DIFFUSERS_PRECISION or set it to bf16.\"\n", + " )\n", + " return FP8_SEPARATE_REPOS[model]\n", + " return MODEL_IDS.get(model, model)\n", + "\n", + "\n", + "\n", + "def materialize_checkpoint_dir(pretrained_source, token) -> Path:\n", + " \"\"\"Return a local checkpoint root (download from Hub when given a repo id).\"\"\"\n", + " path = Path(str(pretrained_source))\n", + " if path.is_dir():\n", + " return path.resolve()\n", + " from huggingface_hub import snapshot_download\n", + "\n", + " return Path(snapshot_download(str(pretrained_source), token=token)).resolve()\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniModularPipeline:\n", " global _pipe, _pipe_model\n", - " model_id = MODEL_IDS.get(model, model)\n", - " if _pipe is not None and _pipe_model == model_id:\n", + " prec = precision_mode()\n", + " pretrained_source = resolve_pretrained_source(model)\n", + " cache_key = f\"{pretrained_source}@{prec}\"\n", + " if _pipe is not None and _pipe_model == cache_key:\n", " return _pipe\n", " release_pipe()\n", - " print(f\"loading {model_id}...\")\n", " t0 = time.time()\n", - " pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)\n", - " pipe.load_components(torch_dtype=torch.bfloat16)\n", + " fp8 = is_fp8_precision()\n", + " token = os.environ.get(\"HF_TOKEN\") or None\n", + " if fp8:\n", + " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", + " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", + " from modelopt.torch.opt import enable_huggingface_checkpointing\n", + "\n", + " enable_huggingface_checkpointing()\n", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", + " pipe = Cosmos3OmniModularPipeline.from_pretrained(\n", + " pretrained_source,\n", + " torch_dtype=torch.bfloat16,\n", + " token=token,\n", + " )\n", + " if fp8:\n", + " # Pre-load the transformer from its full path so ModelOpt 0.44 finds\n", + " # transformer/modelopt_state.pth. Modular load_components keeps subfolder\n", + " # separate, which ModelOpt ignores.\n", + " from diffusers import AutoModel\n", + "\n", + " ckpt_dir = materialize_checkpoint_dir(pretrained_source, token)\n", + " transformer_dir = ckpt_dir / \"transformer\"\n", + " if not (transformer_dir / \"modelopt_state.pth\").is_file():\n", + " raise RuntimeError(\n", + " f\"FP8 requested but no ModelOpt state at {transformer_dir / 'modelopt_state.pth'}. \"\n", + " \"Point COSMOS3_FP8_MODEL_PATH at a pre-quantized checkpoint directory.\"\n", + " )\n", + " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", + " pipe.update_components(transformer=transformer)\n", + " pipe.load_components(pretrained_model_name_or_path=str(ckpt_dir), torch_dtype=torch.bfloat16)\n", + " else:\n", + " pipe.load_components(torch_dtype=torch.bfloat16)\n", " pipe.to(\"cuda\")\n", - " _pipe, _pipe_model = pipe, model_id\n", - " print(f\"loaded pipeline in {time.time() - t0:.1f}s; cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", + " _pipe, _pipe_model = pipe, cache_key\n", + " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", " return _pipe\n", "\n", "\n", @@ -496,7 +578,37 @@ " print(f\"generated in {time.time() - t0:.1f}s\")\n", " export_to_video(videos, str(output_path), fps=spec[\"fps\"], macro_block_size=1)\n", " print(f\"wrote {output_path}\")\n", - " return output_path" + " return output_path\n" + ] + }, + { + "cell_type": "markdown", + "id": "fp8-quantized-checkpoints", + "metadata": {}, + "source": [ + "## Quantized Checkpoints (FP8)\n", + "\n", + "Default loads are BF16 (`COSMOS3_DIFFUSERS_PRECISION=bf16`). Set the env var to `fp8` before running the helper cell's `get_pipe` (or re-run that cell after changing it). Use-case request cells stay the same.\n", + "\n", + "Nano and Super load `nvidia/Cosmos3-Nano-FP8` and `nvidia/Cosmos3-Super-FP8`. Do not pass `revision=fp8` on the BF16 cards.\n", + "\n", + "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", + "\n", + "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory. Transfer preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`.\n", + "\n", + "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + ] + }, + { + "cell_type": "code", + "id": "fp8-precision-toggle", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": [ + "# Default is bf16. Uncomment to load FP8 checkpoints in the cells below.\n", + "# os.environ[\"COSMOS3_DIFFUSERS_PRECISION\"] = \"fp8\"\n", + "print(\"COSMOS3_DIFFUSERS_PRECISION:\", os.environ.get(\"COSMOS3_DIFFUSERS_PRECISION\", \"bf16\"))\n" ] }, { From 9ed508b1de7917717e16f2a7baf2d8c265ccbc80 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Tue, 29 Sep 2026 09:37:49 -0700 Subject: [PATCH 6/7] cookbook(diffusers): fix pose_utils import in ID notebook --- cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb index 647c91fb..29942155 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb @@ -462,7 +462,7 @@ "if str(COSMOS3_REPO) not in sys.path:\n", " sys.path.insert(0, str(COSMOS3_REPO))\n", "from cosmos_framework.data.generator.action.action_normalization import denormalize_action\n", - "from cosmos_framework.data.generator.action.pose_utils import pose_rel_to_abs\n", + "from cosmos_framework.data.generator.action.utils.pose_utils import pose_rel_to_abs\n", "\n", "# frustum: apex + image-rectangle corners (camera +Z forward), and their edges\n", "_FRUSTUM = np.array([[0, 0, 0], [-1, -1, 1], [1, -1, 1], [1, 1, 1], [-1, 1, 1]], float)\n", From dd7088e74d3aa9aa5a9dd43d60dbf9c358859409 Mon Sep 17 00:00:00 2001 From: Yiming Zhao Date: Mon, 19 Oct 2026 11:41:45 -0700 Subject: [PATCH 7/7] cookbook(diffusers): restore FP8 state from fresh Hub cache --- .../action/run_fd_with_diffusers.ipynb | 49 +++++++++++++-- .../action/run_id_with_diffusers.ipynb | 49 +++++++++++++-- .../audiovisual/run_with_diffusers.ipynb | 63 ++++++++++++++----- .../run_video_transfer_with_diffusers.ipynb | 22 ++++++- 4 files changed, 157 insertions(+), 26 deletions(-) diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb index ebc76b2c..c630a954 100644 --- a/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb @@ -609,6 +609,41 @@ " return MODEL_IDS.get(model, model)\n", "\n", "\n", + "def ensure_modelopt_state(pretrained_source, *, token=None, revision=None) -> Path:\n", + " \"\"\"Ensure ModelOpt state is present before Diffusers resolves its Hub allowlist.\"\"\"\n", + " checkpoint_root = Path(str(pretrained_source))\n", + " if isinstance(pretrained_source, Path) or checkpoint_root.is_absolute():\n", + " state_path = checkpoint_root / \"transformer\" / \"modelopt_state.pth\"\n", + " else:\n", + " from huggingface_hub import hf_hub_download\n", + "\n", + " state_path = Path(\n", + " hf_hub_download(\n", + " repo_id=str(pretrained_source),\n", + " filename=\"transformer/modelopt_state.pth\",\n", + " revision=revision,\n", + " token=token,\n", + " )\n", + " )\n", + " if not state_path.is_file():\n", + " raise RuntimeError(\n", + " f\"FP8 requested but no ModelOpt state at {state_path}. \"\n", + " \"Use a ModelOpt FP8 checkpoint or set COSMOS3_FP8_MODEL_PATH to one.\"\n", + " )\n", + " return state_path\n", + "\n", + "\n", + "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", + " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n", + " fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters())\n", + " quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n", + " meta = sum(tensor.is_meta for _, tensor in tensors)\n", + " if not fp8 or not quantizers or meta:\n", + " raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + " return fp8, quantizers, meta\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniPipeline:\n", " global _pipe, _pipe_model\n", " prec = precision_mode()\n", @@ -619,12 +654,16 @@ " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", " t0 = time.time()\n", - " if is_fp8_precision():\n", + " fp8 = is_fp8_precision()\n", + " token = os.environ.get(\"HF_TOKEN\") or None\n", + " if fp8:\n", " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", " from modelopt.torch.opt import enable_huggingface_checkpointing\n", "\n", " enable_huggingface_checkpointing()\n", + " state_path = ensure_modelopt_state(pretrained_source, token=token)\n", + " print(f\"ModelOpt state: {state_path}\")\n", " print(f\"loading FP8 {pretrained_source}...\")\n", " else:\n", " print(f\"loading {pretrained_source}...\")\n", @@ -633,8 +672,11 @@ " torch_dtype=torch.bfloat16,\n", " safety_checker=None,\n", " enable_safety_checker=GUARDRAILS,\n", - " token=os.environ.get(\"HF_TOKEN\") or None,\n", + " token=token,\n", " )\n", + " if fp8:\n", + " fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer)\n", + " print(f\"FP8 verified: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\n", " _pipe, _pipe_model = pipe, cache_key\n", " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", @@ -810,8 +852,7 @@ "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", "\n", "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory.\n", - "\n", - "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + "\n" ] }, { diff --git a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb index 29942155..525ab37e 100644 --- a/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/action/run_id_with_diffusers.ipynb @@ -714,6 +714,41 @@ " return MODEL_IDS.get(model, model)\n", "\n", "\n", + "def ensure_modelopt_state(pretrained_source, *, token=None, revision=None) -> Path:\n", + " \"\"\"Ensure ModelOpt state is present before Diffusers resolves its Hub allowlist.\"\"\"\n", + " checkpoint_root = Path(str(pretrained_source))\n", + " if isinstance(pretrained_source, Path) or checkpoint_root.is_absolute():\n", + " state_path = checkpoint_root / \"transformer\" / \"modelopt_state.pth\"\n", + " else:\n", + " from huggingface_hub import hf_hub_download\n", + "\n", + " state_path = Path(\n", + " hf_hub_download(\n", + " repo_id=str(pretrained_source),\n", + " filename=\"transformer/modelopt_state.pth\",\n", + " revision=revision,\n", + " token=token,\n", + " )\n", + " )\n", + " if not state_path.is_file():\n", + " raise RuntimeError(\n", + " f\"FP8 requested but no ModelOpt state at {state_path}. \"\n", + " \"Use a ModelOpt FP8 checkpoint or set COSMOS3_FP8_MODEL_PATH to one.\"\n", + " )\n", + " return state_path\n", + "\n", + "\n", + "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", + " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n", + " fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters())\n", + " quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n", + " meta = sum(tensor.is_meta for _, tensor in tensors)\n", + " if not fp8 or not quantizers or meta:\n", + " raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + " return fp8, quantizers, meta\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniPipeline:\n", " global _pipe, _pipe_model\n", " prec = precision_mode()\n", @@ -724,12 +759,16 @@ " release_pipe()\n", " diffusers_logging.set_verbosity_info()\n", " t0 = time.time()\n", - " if is_fp8_precision():\n", + " fp8 = is_fp8_precision()\n", + " token = os.environ.get(\"HF_TOKEN\") or None\n", + " if fp8:\n", " # ModelOpt FP8 graph restore + kernels (required before from_pretrained).\n", " import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401\n", " from modelopt.torch.opt import enable_huggingface_checkpointing\n", "\n", " enable_huggingface_checkpointing()\n", + " state_path = ensure_modelopt_state(pretrained_source, token=token)\n", + " print(f\"ModelOpt state: {state_path}\")\n", " print(f\"loading FP8 {pretrained_source}...\")\n", " else:\n", " print(f\"loading {pretrained_source}...\")\n", @@ -738,8 +777,11 @@ " torch_dtype=torch.bfloat16,\n", " safety_checker=None,\n", " enable_safety_checker=GUARDRAILS,\n", - " token=os.environ.get(\"HF_TOKEN\") or None,\n", + " token=token,\n", " )\n", + " if fp8:\n", + " fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer)\n", + " print(f\"FP8 verified: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\n", " _pipe, _pipe_model = pipe, cache_key\n", " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", @@ -977,8 +1019,7 @@ "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", "\n", "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory.\n", - "\n", - "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + "\n" ] }, { diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 3e81f566..1a9af7f9 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -602,7 +602,6 @@ " else:\n", " print(f\"image: {vision_display_path.relative_to(COSMOS_ROOT)}\")\n", " display(Image(filename=str(vision_display_path), width=420))\n", - " print(json.dumps({k: payload[k] for k in [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]}, indent=2))\n", " summary_keys = [\"model_mode\", \"name\", \"enable_sound\", \"num_steps\", \"guidance\", \"shift\", \"fps\", \"num_frames\", \"resolution\", \"aspect_ratio\", \"seed\"]\n", " print(json.dumps({k: payload[k] for k in summary_keys if k in payload}, indent=2))\n", " if payload.get(\"distilled\"):\n", @@ -717,25 +716,50 @@ " ).resolve()\n", "\n", "\n", - "def distilled_components_root(checkpoint_dir: Path) -> str:\n", - " \"\"\"Local Diffusers root for modular `load_components` after transformer preload.\"\"\"\n", - " return str(Path(checkpoint_dir).resolve())\n", + "\n", + "def ensure_modelopt_state(pretrained_source, *, token=None, revision=None) -> Path:\n", + " \"\"\"Ensure ModelOpt state is present before Diffusers resolves its Hub allowlist.\"\"\"\n", + " checkpoint_root = Path(str(pretrained_source))\n", + " if isinstance(pretrained_source, Path) or checkpoint_root.is_absolute():\n", + " state_path = checkpoint_root / \"transformer\" / \"modelopt_state.pth\"\n", + " else:\n", + " from huggingface_hub import hf_hub_download\n", + "\n", + " state_path = Path(\n", + " hf_hub_download(\n", + " repo_id=str(pretrained_source),\n", + " filename=\"transformer/modelopt_state.pth\",\n", + " revision=revision,\n", + " token=token,\n", + " )\n", + " )\n", + " if not state_path.is_file():\n", + " raise RuntimeError(\n", + " f\"FP8 requested but no ModelOpt state at {state_path}. \"\n", + " \"Use a ModelOpt FP8 checkpoint or set COSMOS3_FP8_MODEL_PATH to one.\"\n", + " )\n", + " return state_path\n", "\n", "\n", - "def verify_fp8(model: str) -> None:\n", - " \"\"\"Confirm ModelOpt restored FP8 weights and quantizers (not just float8 tensors).\"\"\"\n", - " transformer = get_pipe(model).transformer\n", + "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", + " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", " tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n", - "\n", - " fp8 = sum(p.dtype == torch.float8_e4m3fn for p in transformer.parameters())\n", + " fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters())\n", " quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n", - " meta = sum(t.is_meta for _, t in tensors)\n", - "\n", - " # Float8 weights alone are insufficient; quantizers prove ModelOpt restored the scales.\n", + " meta = sum(tensor.is_meta for _, tensor in tensors)\n", " if not fp8 or not quantizers or meta:\n", " raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + " return fp8, quantizers, meta\n", + "\n", + "def distilled_components_root(checkpoint_dir: Path) -> str:\n", + " \"\"\"Local Diffusers root for modular `load_components` after transformer preload.\"\"\"\n", + " return str(Path(checkpoint_dir).resolve())\n", + "\n", "\n", - " print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta=0\")\n", + "def verify_fp8(model: str) -> None:\n", + " \"\"\"Confirm ModelOpt restored FP8 weights and quantizers (not just float8 tensors).\"\"\"\n", + " fp8, quantizers, meta = validate_fp8_transformer(get_pipe(model).transformer)\n", + " print(f\"FP8 verified: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", "\n", "\n", "def cuda_allocated_gib() -> float:\n", @@ -832,6 +856,13 @@ " pipe.load_components(torch_dtype=torch.bfloat16)\n", " pipe.enable_safety_checker()\n", " else:\n", + " if fp8:\n", + " state_path = ensure_modelopt_state(\n", + " pretrained_source,\n", + " token=from_kwargs.get(\"token\"),\n", + " revision=from_kwargs.get(\"revision\"),\n", + " )\n", + " print(f\"ModelOpt state: {state_path}\")\n", " pipe = Cosmos3OmniPipeline.from_pretrained(\n", " pretrained_source,\n", " torch_dtype=torch.bfloat16,\n", @@ -840,6 +871,9 @@ " **from_kwargs,\n", " )\n", " pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FIXED_SAMPLING[\"shift\"])\n", + " if fp8:\n", + " fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer)\n", + " print(f\"FP8 verified: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\n", " _pipe = pipe\n", " _pipe_model = cache_key\n", @@ -994,8 +1028,7 @@ "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). Distilled 4-Step stays all W8A8. See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", "\n", "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory. Distilled FP8 preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`. After load, `verify_fp8(\"Cosmos3-Nano\")` (or another supported model) checks FP8 weights and quantizer modules.\n", - "\n", - "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + "\n" ] }, { diff --git a/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb index 027bb764..22d65c96 100644 --- a/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_diffusers.ipynb @@ -484,6 +484,17 @@ " return Path(snapshot_download(str(pretrained_source), token=token)).resolve()\n", "\n", "\n", + "def validate_fp8_transformer(transformer) -> tuple[int, int, int]:\n", + " \"\"\"Reject FP8 weights loaded without their ModelOpt quantizer graph.\"\"\"\n", + " tensors = list(transformer.named_parameters()) + list(transformer.named_buffers())\n", + " fp8 = sum(parameter.dtype == torch.float8_e4m3fn for parameter in transformer.parameters())\n", + " quantizers = sum(\"quantizer\" in name.lower() for name, _ in transformer.named_modules())\n", + " meta = sum(tensor.is_meta for _, tensor in tensors)\n", + " if not fp8 or not quantizers or meta:\n", + " raise RuntimeError(f\"Invalid FP8 load: weights={fp8}, quantizers={quantizers}, meta={meta}\")\n", + " return fp8, quantizers, meta\n", + "\n", + "\n", "def get_pipe(model: str) -> Cosmos3OmniModularPipeline:\n", " global _pipe, _pipe_model\n", " prec = precision_mode()\n", @@ -524,9 +535,15 @@ " )\n", " transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)\n", " pipe.update_components(transformer=transformer)\n", - " pipe.load_components(pretrained_model_name_or_path=str(ckpt_dir), torch_dtype=torch.bfloat16)\n", + " # Keep the remaining component sources from modular_model_index.json. The FP8 repo's\n", + " # root modelopt_state.pth belongs to the Reasoner and must not be applied to the VAE or\n", + " # audio tokenizer.\n", + " pipe.load_components(torch_dtype=torch.bfloat16, token=token)\n", " else:\n", " pipe.load_components(torch_dtype=torch.bfloat16)\n", + " if fp8:\n", + " fp8_weights, quantizers, meta = validate_fp8_transformer(pipe.transformer)\n", + " print(f\"FP8 verified: weights={fp8_weights}, quantizers={quantizers}, meta={meta}\")\n", " pipe.to(\"cuda\")\n", " _pipe, _pipe_model = pipe, cache_key\n", " print(f\"loaded pipeline in {time.time() - t0:.1f}s (precision={prec}); cuda allocated {cuda_allocated_gib():.1f} GiB\")\n", @@ -595,8 +612,7 @@ "Diffusers follows the checkpoint schedule (W8A16 on the first and last few steps, W8A8 in between). See the [Cosmos 3 Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3).\n", "\n", "Optional: point `COSMOS3_FP8_MODEL_PATH` at a local Diffusers checkpoint directory. Transfer preloads `transformer/` so ModelOpt 0.44 restores `modelopt_state.pth`.\n", - "\n", - "The dedicated `*-FP8` Hub repos may stay private until the official model-card release.\n" + "\n" ] }, {