docs: fix Diffusers FP8 modular restore

README.md

@@ -514,6 +514,7 @@ uv pip install \
   huggingface_hub \
   imageio \
   imageio-ffmpeg \
+  nvidia-modelopt==0.44.0 \
   torch \
   torchvision \
   transformers
@@ -523,18 +524,35 @@ The safety checker requires access to the gated [nvidia/Cosmos-1.0-Guardrail](ht
 
 ### Example: Text to Image Generation
 
+Enable ModelOpt checkpoint restoration before loading the FP8 transformer. ModelOpt 0.44 does not use a Hub `subfolder` when locating `modelopt_state.pth`, so the example materializes the `fp8` revision, loads `transformer/` from its full local path, and pins it before loading the remaining modular components.
+
 ```python
 import json
+from pathlib import Path
 
+import modelopt.torch.quantization.backends.fp8_per_tensor_gemm  # noqa: F401
 import torch
-from diffusers import Cosmos3DistilledModularPipeline
+from diffusers import AutoModel, Cosmos3DistilledModularPipeline
+from huggingface_hub import snapshot_download
+from modelopt.torch.opt import enable_huggingface_checkpointing
 
 json_prompt = json.load(open("assets/example_caption.json"))
 json_prompt["resolution"] = {"H": 768, "W": 768}
 json_prompt["aspect_ratio"] = "1,1"
 
-pipe = Cosmos3DistilledModularPipeline.from_pretrained("nvidia/Cosmos3-Super-Text2Image-4Step")
-pipe.load_components(torch_dtype=torch.bfloat16)
+model_id = "nvidia/Cosmos3-Super-Text2Image-4Step"
+revision = "fp8"
+enable_huggingface_checkpointing()
+
+pipe = Cosmos3DistilledModularPipeline.from_pretrained(model_id, revision=revision)
+checkpoint_dir = Path(snapshot_download(model_id, revision=revision)).resolve()
+transformer_dir = checkpoint_dir / "transformer"
+modelopt_state = transformer_dir / "modelopt_state.pth"
+if not modelopt_state.is_file():
+    raise RuntimeError(f"FP8 requested but no ModelOpt state at {modelopt_state}")
+transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)
+pipe.update_components(transformer=transformer)
+pipe.load_components(pretrained_model_name_or_path=str(checkpoint_dir), torch_dtype=torch.bfloat16)
 pipe.enable_safety_checker()
 pipe.to("cuda")