diff --git a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index 2c058af6..d83734e9 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", @@ -643,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", @@ -680,12 +752,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 +788,44 @@ " \"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", + " 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", - " 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 +954,47 @@ " 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", + "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" + ] + }, + { + "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",