[Cosmos3] Mixed W8A8/W8A16 denoising for ModelOpt FP8 checkpoints

docs/source/en/api/pipelines/cosmos3.md

@@ -43,6 +43,33 @@ Two checkpoints are released on the Hub — [`nvidia/Cosmos3-Nano`](https://hugg
 > [!TIP]
 > Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reusing-models-in-multiple-pipelines) section to learn how to efficiently load the same components into multiple pipelines.
 
+## FP8 mixed W8A8/W8A16 denoising
+
+Official ModelOpt FP8 checkpoints live on the Hub `fp8` revision (for example [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) with `revision="fp8"`). The serialized weights are static W8A8. Video Nano / Super / Super-I2V checkpoints also store a `quantization_config.runtime.diffusion_step_policy` on the transformer: the **first 3 and last 3** denoising steps run **W8A16** (dequantized FP8 weights, `torch.nn.functional.linear`), and the middle steps keep native **W8A8**. Precision is chosen once per scheduler step so CFG cond/uncond calls match. Distilled 4-step and Super-T2I FP8 checkpoints omit that policy (`runtime` is `null`) and stay native W8A8 on every step.
+
+Loading those weights still uses [`NVIDIAModelOptConfig`](../../quantization/modelopt) as in the ModelOpt guide. After restore, mixed precision is **on by default** when the checkpoint declares the policy — you do not pass a format flag:
+
+```python
+import torch
+from diffusers import Cosmos3OmniPipeline
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Nano",
+    revision="fp8",
+    dtype=torch.bfloat16,
+    device_map="cuda",
+)
+result = pipe(prompt="...", num_inference_steps=35)  # 3×W8A16 / 29×W8A8 / 3×W8A16
+```
+
+Call-site overrides:
+
+- `mixed_precision_format="none"` disables the schedule only (quantized W8A8 remains).
+- `mixed_precision_format="fp8"` forces the first/last-N schedule even if the checkpoint has no policy.
+- `mixed_precision_first_steps` / `mixed_precision_last_steps` / `mixed_precision_reasoner_policy` override the checkpoint counts and reasoner path (`"high_precision"` = W8A16, `"base_precision"` = native W8A8).
+
+These kwargs are not Accelerate `mixed_precision`. They only select W8A8 vs W8A16 on Cosmos3 ModelOpt linears.
+
 ## Prompt upsampling
 
 Cosmos 3 was trained on long, highly descriptive captions. For optimal quality, short text prompts should be **upsampled into a specific JSON structure** before they are passed to the pipeline. The upsampler lives in the [cosmos-framework](https://github.com/NVIDIA/cosmos-framework) package.
@@ -1117,6 +1144,9 @@ config (from the checkpoint's `modular_model_index.json`) and `guidance_scale` i
 1.0 since guidance is baked into the weights — passing any other value for either raises an error,
 and `negative_prompt` is warned about and ignored.
 
+FP8 distilled checkpoints (`revision="fp8"`) do not declare a mixed-precision policy, so every
+step stays native W8A8.
+
 Prompts follow the same descriptive JSON structure as the non-distilled models, so short text
 must be upsampled first — use `--mode text2image` (T2I) or `--mode image2video` (I2V) as
 described in [Prompt upsampling](#prompt-upsampling), then pass the JSON via `json.dumps(...)`.

src/diffusers/modular_pipelines/cosmos/denoise.py

@@ -3,6 +3,11 @@
 import torch
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+from ...pipelines.cosmos.mixed_precision import (
+    Cosmos3MixedPrecisionConfig,
+    apply_cosmos3_mixed_precision_step,
+    reset_cosmos3_mixed_precision,
+)
 from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ..modular_pipeline import (
     BlockState,
@@ -463,18 +468,61 @@ def loop_inputs(self) -> list[InputParam]:
             InputParam(
                 name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps."
             ),
+            InputParam(
+                name="mixed_precision_format",
+                type_hint=str,
+                default=None,
+                description="None reads the checkpoint diffusion_step_policy; 'fp8' forces mixed precision; 'none' disables it.",
+            ),
+            InputParam(
+                name="mixed_precision_first_steps",
+                type_hint=int,
+                default=None,
+                description="Optional override for the leading W8A16 step count.",
+            ),
+            InputParam(
+                name="mixed_precision_last_steps",
+                type_hint=int,
+                default=None,
+                description="Optional override for the trailing W8A16 step count.",
+            ),
+            InputParam(
+                name="mixed_precision_reasoner_policy",
+                type_hint=str,
+                default=None,
+                description="Optional override: W8A16 ('high_precision') or native W8A8 ('base_precision') for the reasoner path.",
+            ),
         ]
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
-        with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
-            for i, t in enumerate(block_state.timesteps):
-                components, block_state = self.loop_step(components, block_state, i=i, t=t)
-                if i == len(block_state.timesteps) - 1 or (
-                    (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
-                ):
-                    progress_bar.update()
+        mixed_precision = Cosmos3MixedPrecisionConfig.resolve(
+            components.transformer,
+            mixed_precision_format=getattr(block_state, "mixed_precision_format", None),
+            mixed_precision_first_steps=getattr(block_state, "mixed_precision_first_steps", None),
+            mixed_precision_last_steps=getattr(block_state, "mixed_precision_last_steps", None),
+            mixed_precision_reasoner_policy=getattr(block_state, "mixed_precision_reasoner_policy", None),
+        )
+        trace = []
+        try:
+            with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+                for i, t in enumerate(block_state.timesteps):
+                    apply_cosmos3_mixed_precision_step(
+                        components.transformer,
+                        mixed_precision,
+                        i,
+                        len(block_state.timesteps),
+                        trace=trace,
+                    )
+                    components, block_state = self.loop_step(components, block_state, i=i, t=t)
+                    if i == len(block_state.timesteps) - 1 or (
+                        (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+                    ):
+                        progress_bar.update()
+        finally:
+            reset_cosmos3_mixed_precision(components.transformer, mixed_precision)
+        components._mixed_precision_trace = trace
         self.set_block_state(state, block_state)
         return components, state