docs: fix Diffusers FP8 modular restore

README.md

@@ -510,6 +510,7 @@ uv pip install \
   huggingface_hub \
   imageio \
   imageio-ffmpeg \
+  nvidia-modelopt==0.44.0 \
   torch \
   torchvision \
   transformers
@@ -519,18 +520,35 @@ The safety checker requires access to the gated [nvidia/Cosmos-1.0-Guardrail](ht
 
 ### Example: Image to Video 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
 from diffusers.utils import export_to_video, load_image
 
 json_prompt = json.load(open("assets/example_prompt.json"))
 image = load_image("assets/example_first_frame.png")
 
-pipe = Cosmos3DistilledModularPipeline.from_pretrained("nvidia/Cosmos3-Super-Image2Video-4Step")
-pipe.load_components(torch_dtype=torch.bfloat16)
+model_id = "nvidia/Cosmos3-Super-Image2Video-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")