ConstBob commited on
Commit
a3d514c
·
verified ·
1 Parent(s): 220a1b2

docs: fix Diffusers FP8 modular restore

Browse files

Load the fp8 revision explicitly, enable ModelOpt checkpoint restoration, and preload the transformer from its full local path before loading the remaining modular components.

Files changed (1) hide show
  1. README.md +21 -3
README.md CHANGED
@@ -510,6 +510,7 @@ uv pip install \
510
  huggingface_hub \
511
  imageio \
512
  imageio-ffmpeg \
 
513
  torch \
514
  torchvision \
515
  transformers
@@ -519,18 +520,35 @@ The safety checker requires access to the gated [nvidia/Cosmos-1.0-Guardrail](ht
519
 
520
  ### Example: Image to Video Generation
521
 
 
 
522
  ```python
523
  import json
 
524
 
 
525
  import torch
526
- from diffusers import Cosmos3DistilledModularPipeline
 
 
527
  from diffusers.utils import export_to_video, load_image
528
 
529
  json_prompt = json.load(open("assets/example_prompt.json"))
530
  image = load_image("assets/example_first_frame.png")
531
 
532
- pipe = Cosmos3DistilledModularPipeline.from_pretrained("nvidia/Cosmos3-Super-Image2Video-4Step")
533
- pipe.load_components(torch_dtype=torch.bfloat16)
 
 
 
 
 
 
 
 
 
 
 
534
  pipe.enable_safety_checker()
535
  pipe.to("cuda")
536
 
 
510
  huggingface_hub \
511
  imageio \
512
  imageio-ffmpeg \
513
+ nvidia-modelopt==0.44.0 \
514
  torch \
515
  torchvision \
516
  transformers
 
520
 
521
  ### Example: Image to Video Generation
522
 
523
+ 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.
524
+
525
  ```python
526
  import json
527
+ from pathlib import Path
528
 
529
+ import modelopt.torch.quantization.backends.fp8_per_tensor_gemm # noqa: F401
530
  import torch
531
+ from diffusers import AutoModel, Cosmos3DistilledModularPipeline
532
+ from huggingface_hub import snapshot_download
533
+ from modelopt.torch.opt import enable_huggingface_checkpointing
534
  from diffusers.utils import export_to_video, load_image
535
 
536
  json_prompt = json.load(open("assets/example_prompt.json"))
537
  image = load_image("assets/example_first_frame.png")
538
 
539
+ model_id = "nvidia/Cosmos3-Super-Image2Video-4Step"
540
+ revision = "fp8"
541
+ enable_huggingface_checkpointing()
542
+
543
+ pipe = Cosmos3DistilledModularPipeline.from_pretrained(model_id, revision=revision)
544
+ checkpoint_dir = Path(snapshot_download(model_id, revision=revision)).resolve()
545
+ transformer_dir = checkpoint_dir / "transformer"
546
+ modelopt_state = transformer_dir / "modelopt_state.pth"
547
+ if not modelopt_state.is_file():
548
+ raise RuntimeError(f"FP8 requested but no ModelOpt state at {modelopt_state}")
549
+ transformer = AutoModel.from_pretrained(str(transformer_dir), torch_dtype=torch.bfloat16)
550
+ pipe.update_components(transformer=transformer)
551
+ pipe.load_components(pretrained_model_name_or_path=str(checkpoint_dir), torch_dtype=torch.bfloat16)
552
  pipe.enable_safety_checker()
553
  pipe.to("cuda")
554