diff --git a/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb b/cookbooks/cosmos3/generator/action/run_fd_with_diffusers.ipynb index 2ae342f8..c630a954 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,107 @@ " 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 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", - " 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", + " 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", " 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", + " 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, 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 +838,35 @@ " 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" + ] + }, + { + "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..525ab37e 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." ] @@ -454,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", @@ -676,25 +684,107 @@ " 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 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", - " 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", + " 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", " 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", + " 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, 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 +1005,35 @@ " )\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" + ] + }, + { + "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/audiovisual/run_with_diffusers.ipynb b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb index adc9cbed..1a9af7f9 100644 --- a/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb +++ b/cookbooks/cosmos3/generator/audiovisual/run_with_diffusers.ipynb @@ -19,7 +19,9 @@ "\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", - "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" + "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", + "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" ] }, { @@ -57,6 +59,7 @@ "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" ] }, @@ -103,6 +106,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 +118,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 +164,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", @@ -596,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", @@ -644,18 +649,119 @@ " \"Cosmos3-Super-Image2Video-4Step\",\n", "}\n", "\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", "\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_model_id(model: str) -> str:\n", " return MODEL_IDS.get(model, model)\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() and model in FP8_SEPARATE_REPOS:\n", + " return FP8_SEPARATE_REPOS[model]\n", + " return resolve_model_id(model)\n", + "\n", + "\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", " 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", + "\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", + "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", + " 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", " return torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0\n", "\n", @@ -692,13 +798,31 @@ "\n", "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", + " 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", + " 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", + " 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", + " print(f\"loading FP8 {pretrained_source}...\")\n", + " else:\n", + " print(f\"loading {pretrained_source}...\")\n", + "\n", " if is_distilled_model(model):\n", " if Cosmos3DistilledModularPipeline is None:\n", " raise RuntimeError(\n", @@ -706,24 +830,54 @@ " \"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", + " 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", - " 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", + " 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 = 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", @@ -859,6 +1013,36 @@ " display(Image(filename=str(src), width=720))\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 / 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" + ] + }, + { + "cell_type": "code", + "id": "fp8-precision-toggle", + "metadata": {}, + "outputs": [], + "execution_count": null, + "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": "d988d45b", @@ -1746,4 +1930,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} 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..22d65c96 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,110 @@ " 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 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", - " 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", + " # 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, 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 +595,36 @@ " 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" + ] + }, + { + "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" ] }, {