Cosmos3 Distilled support (#14177)

* Cosmos3 distilled support (dedicated modular pipeline + blocks)

* Address review comments

* Refactor unit test to use ModularPipelineTesterMixin

* Docs fix

* Tests refactor

* Pass generator to SDE scheduler  for reproducibility

* Apply suggestion from @yiyixuxu

* Distilled text encoder step bugfix

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>

Author: yzhautouskay

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

@@ -989,11 +989,79 @@ videos = pipe(
 export_to_video(videos, "cosmos3_modular_transfer_edge.mp4", fps=30, macro_block_size=1)
 ```
 
+### Distilled (few-step) text-to-image and image-to-video
+
+Few-step distilled checkpoints are served by [`Cosmos3DistilledModularPipeline`] (blocks:
+`Cosmos3DistilledBlocks`); the base [`Cosmos3OmniModularPipeline`] and [`Cosmos3OmniPipeline`] do
+not support them. `num_inference_steps` is fixed to the length of the `distilled_sigmas` pipeline
+config (from the checkpoint's `modular_model_index.json`) and `guidance_scale` is forced to
+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.
+
+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(...)`.
+
+```python
+import json
+import torch
+from diffusers import Cosmos3DistilledModularPipeline
+from diffusers.utils import export_to_video, load_image
+
+# JSON-upsampled prompt (see "Prompt upsampling" above).
+json_prompt = json.load(open("assets/example_t2i_prompt.json"))
+
+repo = "nvidia/Cosmos3-Super-Text2Image-4Step"
+pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+# text-to-image (distilled)
+videos = pipe(
+    prompt=json.dumps(json_prompt),
+    num_frames=1,
+    height=720,
+    width=1280,
+    output="videos",
+)
+videos[0].save("cosmos3_distilled_t2i.jpg", format="JPEG", quality=85)
+
+# image-to-video (distilled) — load the I2V repo instead
+# JSON-upsampled prompt (see "Prompt upsampling" above); upsampled from the source prompt
+# "The right robotic hand picks up the red sphere on the shelf."
+json_prompt_i2v = json.load(open("assets/example_i2v_prompt.json"))
+
+repo_i2v = "nvidia/Cosmos3-Super-Image2Video-4Step"
+pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo_i2v, torch_dtype=torch.bfloat16)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+image = load_image(
+    "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg"
+)
+videos = pipe(
+    prompt=json.dumps(json_prompt_i2v),
+    image=image,
+    num_frames=189,
+    height=720,
+    width=1280,
+    output="videos",
+)
+export_to_video(videos, "cosmos3_distilled_i2v.mp4", fps=24, macro_block_size=1)
+```
+
 [[autodoc]] Cosmos3OmniModularPipeline
 
 - all
 - __call__
 
+## Cosmos3DistilledModularPipeline
+
+[[autodoc]] Cosmos3DistilledModularPipeline
+
+- all
+- __call__
+
 ## CosmosActionCondition
 
 [[autodoc]] CosmosActionCondition

src/diffusers/__init__.py

@@ -489,6 +489,8 @@
         [
             "AnimaAutoBlocks",
             "AnimaModularPipeline",
+            "Cosmos3DistilledBlocks",
+            "Cosmos3DistilledModularPipeline",
             "Cosmos3OmniBlocks",
             "Cosmos3OmniModularPipeline",
             "ErnieImageAutoBlocks",
@@ -1349,6 +1351,8 @@
         from .modular_pipelines import (
             AnimaAutoBlocks,
             AnimaModularPipeline,
+            Cosmos3DistilledBlocks,
+            Cosmos3DistilledModularPipeline,
             Cosmos3OmniBlocks,
             Cosmos3OmniModularPipeline,
             ErnieImageAutoBlocks,

src/diffusers/modular_pipelines/__init__.py

@@ -98,6 +98,8 @@
         "AnimaModularPipeline",
     ]
     _import_structure["cosmos"] = [
+        "Cosmos3DistilledBlocks",
+        "Cosmos3DistilledModularPipeline",
         "Cosmos3OmniBlocks",
         "Cosmos3OmniModularPipeline",
     ]
@@ -128,7 +130,12 @@
     else:
         from .anima import AnimaAutoBlocks, AnimaModularPipeline
         from .components_manager import ComponentsManager
-        from .cosmos import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline
+        from .cosmos import (
+            Cosmos3DistilledBlocks,
+            Cosmos3DistilledModularPipeline,
+            Cosmos3OmniBlocks,
+            Cosmos3OmniModularPipeline,
+        )
         from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline
         from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
         from .flux2 import (

src/diffusers/modular_pipelines/cosmos/__init__.py

@@ -22,7 +22,8 @@
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
     _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"]
-    _import_structure["modular_pipeline"] = ["Cosmos3OmniModularPipeline"]
+    _import_structure["modular_blocks_cosmos3_distilled"] = ["Cosmos3DistilledBlocks"]
+    _import_structure["modular_pipeline"] = ["Cosmos3DistilledModularPipeline", "Cosmos3OmniModularPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -32,7 +33,8 @@
         from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
         from .modular_blocks_cosmos3 import Cosmos3OmniBlocks
-        from .modular_pipeline import Cosmos3OmniModularPipeline
+        from .modular_blocks_cosmos3_distilled import Cosmos3DistilledBlocks
+        from .modular_pipeline import Cosmos3DistilledModularPipeline, Cosmos3OmniModularPipeline
 else:
     import sys
 

src/diffusers/modular_pipelines/cosmos/before_denoise.py

@@ -5,7 +5,7 @@
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
 from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
-from ...schedulers import UniPCMultistepScheduler
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ...utils.torch_utils import randn_tensor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
@@ -116,6 +116,11 @@ def intermediate_outputs(self) -> list[OutputParam]:
                 type_hint=list[int],
                 description="Indexes of conditioned vision latent frames.",
             ),
+            OutputParam(
+                "vision_conditioning_latents",
+                type_hint=torch.Tensor,
+                description="Clean encoded vision latents used to re-anchor image conditioning each step.",
+            ),
         ]
 
     @torch.no_grad()
@@ -166,6 +171,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
             block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False
         ).flatten()
         block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes]
+        block_state.vision_conditioning_latents = x0_tokens_vision
 
         self.set_block_state(state, block_state)
         return components, state
@@ -1237,3 +1243,89 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
         )
         self.set_block_state(state, block_state)
         return components, state
+
+
+class Cosmos3DistilledSetTimestepsStep(ModularPipelineBlocks):
+    model_name = "cosmos3-omni"
+
+    @property
+    def description(self) -> str:
+        return "Initializes the fixed distilled sampling schedule from the pipeline's `distilled_sigmas` config."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [
+            ConfigSpec(name="is_distilled", default=True),
+            ConfigSpec(name="distilled_sigmas", default=None),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_inference_steps", required=False, default=None),
+            InputParam(
+                name="guidance_scale",
+                type_hint=float,
+                default=None,
+                description=(
+                    "Unused for distilled checkpoints; classifier-free guidance is baked into the weights and the "
+                    "scale is forced to 1.0. Passing a value other than 1.0 raises an error."
+                ),
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."),
+            OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."),
+            OutputParam(
+                "num_inference_steps",
+                type_hint=int,
+                description="Resolved number of denoising steps (fixed by the distilled schedule).",
+            ),
+            OutputParam(
+                name="guidance_scale",
+                type_hint=float,
+                description="Resolved classifier-free guidance scale (always 1.0 for distilled checkpoints).",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        device = components._execution_device
+
+        sigmas = components.config.distilled_sigmas
+        if not sigmas:
+            raise ValueError(
+                "Cosmos3DistilledSetTimestepsStep requires the pipeline config `distilled_sigmas` to be set "
+                "(populated from the distilled checkpoint's `modular_model_index.json`). Load a distilled Cosmos3 "
+                "checkpoint or use `Cosmos3OmniModularPipeline` for base checkpoints."
+            )
+        sigmas = [float(s) for s in sigmas]
+        distilled_steps = len(sigmas)
+
+        if block_state.num_inference_steps is not None and block_state.num_inference_steps != distilled_steps:
+            raise ValueError(
+                "This is a distilled checkpoint; the step count is fixed by the pipeline's "
+                f"`distilled_sigmas` config ({distilled_steps} steps). "
+                f"`num_inference_steps` must be {distilled_steps} or left unset (got {block_state.num_inference_steps})."
+            )
+        if block_state.guidance_scale is not None and block_state.guidance_scale != 1.0:
+            raise ValueError(
+                "This is a distilled checkpoint; classifier-free guidance is baked into the weights. "
+                f"`guidance_scale` must be 1.0 or left unset (got {block_state.guidance_scale})."
+            )
+
+        components.scheduler.set_timesteps(sigmas=sigmas, device=device)
+        block_state.num_inference_steps = distilled_steps
+        block_state.guidance_scale = 1.0
+        block_state.timesteps = components.scheduler.timesteps
+        block_state.num_warmup_steps = len(block_state.timesteps) - distilled_steps * components.scheduler.order
+        self.set_block_state(state, block_state)
+        return components, state

src/diffusers/modular_pipelines/cosmos/denoise.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
-from ...schedulers import UniPCMultistepScheduler
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ..modular_pipeline import (
     BlockState,
     LoopSequentialPipelineBlocks,
@@ -282,6 +282,71 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
         return components, block_state
 
 
+class Cosmos3DistilledVisionLoopSchedulerStep(ModularPipelineBlocks):
+    model_name = "cosmos3-omni"
+
+    @property
+    def description(self) -> str:
+        return "Updates vision latents after one distilled denoising iteration, re-anchoring conditioned frames."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("latents", required=True, description="Noisy vision latents to update."),
+            InputParam(
+                name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity."
+            ),
+            InputParam(
+                name="vision_condition_mask",
+                type_hint=torch.Tensor,
+                required=True,
+                description="Mask marking conditioned vision latent frames.",
+            ),
+            InputParam(
+                name="vision_conditioning_latents",
+                type_hint=torch.Tensor,
+                default=None,
+                description="Clean encoded vision latents for re-anchoring conditioned frames.",
+            ),
+            InputParam(
+                name="vision_condition_indexes_for_pack",
+                type_hint=list,
+                default=None,
+                description="Indexes of conditioned vision latent frames; non-empty for image-to-video.",
+            ),
+            InputParam.template("generator"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam.template("latents")]
+
+    @torch.no_grad()
+    def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        # Pass the generator so the scheduler's stochastic (SDE) re-noising is seedable/reproducible.
+        block_state.latents = components.scheduler.step(
+            block_state.velocity_vision.unsqueeze(0),
+            t,
+            block_state.latents.unsqueeze(0),
+            generator=block_state.generator,
+            return_dict=False,
+        )[0].squeeze(0)
+
+        # Distilled checkpoints use stochastic (SDE) scheduler steps that re-noise every position.
+        # Re-anchor conditioned frames to the clean encoded reference after each step.
+        has_image_condition = bool(block_state.vision_condition_indexes_for_pack)
+        if has_image_condition and block_state.vision_conditioning_latents is not None:
+            mask = block_state.vision_condition_mask
+            reference = block_state.vision_conditioning_latents.to(block_state.latents.dtype)
+            block_state.latents = mask * reference + (1.0 - mask) * block_state.latents
+
+        return components, block_state
+
+
 class Cosmos3SoundLoopSchedulerStep(ModularPipelineBlocks):
     model_name = "cosmos3-omni"
 
@@ -427,6 +492,27 @@ def description(self) -> str:
         return "Runs the vision-only Cosmos3 denoising loop."
 
 
+class Cosmos3DistilledVisionDenoiseStep(Cosmos3DenoiseLoopWrapper):
+    model_name = "cosmos3-omni"
+    block_classes = [
+        Cosmos3VisionLoopPrepareStep,
+        Cosmos3LoopDenoiser,
+        Cosmos3DistilledVisionLoopSchedulerStep,
+    ]
+    block_names = ["prepare_vision", "denoiser", "update_vision"]
+
+    @property
+    def description(self) -> str:
+        return "Runs the vision-only distilled Cosmos3 denoising loop."
+
+    @property
+    def loop_expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+            ComponentSpec("transformer", Cosmos3OmniTransformer),
+        ]
+
+
 class Cosmos3VisionSoundDenoiseStep(Cosmos3DenoiseLoopWrapper):
     block_classes = [
         Cosmos3VisionLoopPrepareStep,