Cosmos3 Distilled support

Author: yzhautouskay

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

@@ -934,6 +934,73 @@ if outputs["action"] is not None:
         json.dump(outputs["action"][0].tolist(), f)
 ```
 
+### Distilled (few-step) text-to-image and image-to-video
+
+Few-step distilled checkpoints ship a fixed sigma schedule in
+`scheduler.config.fixed_step_sampler_config.t_list` and bake classifier-free guidance into
+the weights. **Only the modular pipeline supports them today** — the task-based
+[`Cosmos3OmniPipeline`] does not implement the distilled contract.
+
+Load a distilled repo and call the modular pipeline **without** `num_inference_steps` or
+`guidance_scale`; `Cosmos3SetTimestepsStep` resolves both from the scheduler config
+(fixed step count, `guidance_scale=1.0`). Because classifier-free guidance is baked into the
+weights, negative prompts are unused.
+
+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 Cosmos3OmniModularPipeline
+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 = Cosmos3OmniModularPipeline.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 = Cosmos3OmniModularPipeline.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)
+```
+
+For a CLI wrapper with warmup / iteration timing, see
+[`examples/cosmos3/inference_cosmos3_modular_distilled.py`](../../../examples/cosmos3/inference_cosmos3_modular_distilled.py).
+
 [[autodoc]] Cosmos3OmniModularPipeline
 
 - all

src/diffusers/modular_pipelines/cosmos/before_denoise.py

@@ -4,7 +4,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 SchedulerMixin
 from ...utils.torch_utils import randn_tensor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
@@ -115,6 +115,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()
@@ -165,6 +170,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
@@ -181,7 +187,7 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("transformer", Cosmos3OmniTransformer),
-            ComponentSpec("scheduler", UniPCMultistepScheduler),
+            ComponentSpec("scheduler", SchedulerMixin),
         ]
 
     @property
@@ -257,7 +263,7 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("transformer", Cosmos3OmniTransformer),
-            ComponentSpec("scheduler", UniPCMultistepScheduler),
+            ComponentSpec("scheduler", SchedulerMixin),
         ]
 
     @property
@@ -934,29 +940,54 @@ def description(self) -> str:
 
     @property
     def expected_components(self) -> list[ComponentSpec]:
-        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+        return [ComponentSpec("scheduler", SchedulerMixin)]
 
     @property
     def inputs(self) -> list[InputParam]:
         return [
-            InputParam.template("num_inference_steps", required=True),
+            InputParam.template("num_inference_steps", required=False),
+            InputParam(
+                name="guidance_scale",
+                type_hint=float,
+                default=None,
+                description="Classifier-free guidance scale. Forced to 1.0 for distilled checkpoints.",
+            ),
         ]
 
     @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.",
+            ),
+            OutputParam(
+                name="guidance_scale",
+                type_hint=float,
+                description="Resolved classifier-free guidance scale for the denoising loop.",
+            ),
         ]
 
     @torch.no_grad()
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
-        block_state.timesteps = components.scheduler.timesteps
-        block_state.num_warmup_steps = (
-            len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order
+        num_inference_steps, guidance_scale = components.resolve_inference_schedule(
+            block_state.num_inference_steps,
+            block_state.guidance_scale,
         )
+        block_state.num_inference_steps = num_inference_steps
+        block_state.guidance_scale = guidance_scale
+
+        if components.is_distilled_checkpoint:
+            components.scheduler.set_timesteps(sigmas=components.distilled_sigmas, device=device)
+        else:
+            components.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        block_state.timesteps = components.scheduler.timesteps
+        block_state.num_warmup_steps = len(block_state.timesteps) - num_inference_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 SchedulerMixin
 from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 from .modular_pipeline import Cosmos3OmniModularPipeline
@@ -254,7 +254,7 @@ def description(self) -> str:
 
     @property
     def expected_components(self) -> list[ComponentSpec]:
-        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+        return [ComponentSpec("scheduler", SchedulerMixin)]
 
     @property
     def inputs(self) -> list[InputParam]:
@@ -263,6 +263,24 @@ def inputs(self) -> list[InputParam]:
             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.",
+            ),
         ]
 
     @property
@@ -274,6 +292,19 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
         block_state.latents = components.scheduler.step(
             block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), 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 (
+            components.is_distilled_checkpoint
+            and 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
 
 
@@ -295,7 +326,7 @@ def inputs(self) -> list[InputParam]:
             ),
             InputParam(
                 name="sound_scheduler",
-                type_hint=UniPCMultistepScheduler,
+                type_hint=SchedulerMixin,
                 required=True,
                 description="Scheduler used to update sound latents.",
             ),
@@ -334,7 +365,7 @@ def inputs(self) -> list[InputParam]:
             ),
             InputParam(
                 name="action_scheduler",
-                type_hint=UniPCMultistepScheduler,
+                type_hint=SchedulerMixin,
                 required=True,
                 description="Scheduler used to update action latents.",
             ),
@@ -381,7 +412,7 @@ def description(self) -> str:
     @property
     def loop_expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("scheduler", UniPCMultistepScheduler),
+            ComponentSpec("scheduler", SchedulerMixin),
             ComponentSpec("transformer", Cosmos3OmniTransformer),
         ]
 

src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py

@@ -242,7 +242,7 @@ class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text-and-vision Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`)
 
       Inputs:
           cond_input_ids (`None`):
@@ -310,7 +310,7 @@ class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, and sound Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`)
 
       Inputs:
           cond_input_ids (`None`):
@@ -391,7 +391,7 @@ class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, and action Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`)
 
       Inputs:
           cond_input_ids (`None`):
@@ -476,7 +476,7 @@ class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks):
     Runs the text, vision, sound, and action Cosmos3 denoising workflow.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`)
 
       Inputs:
           cond_input_ids (`None`):
@@ -576,7 +576,7 @@ class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks):
        - vision runs otherwise.
 
       Components:
-          transformer (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`)
+          transformer (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`)
 
       Inputs:
           cond_input_ids (`None`):
@@ -690,7 +690,7 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
 
       Components:
           text_tokenizer (`AutoTokenizer`) video_processor (`VideoProcessor`) vae (`AutoencoderKLWan`) transformer
-          (`Cosmos3OmniTransformer`) scheduler (`UniPCMultistepScheduler`) sound_tokenizer
+          (`Cosmos3OmniTransformer`) scheduler (`SchedulerMixin`) sound_tokenizer
           (`Cosmos3AVAEAudioTokenizer`)
 
       Inputs:

src/diffusers/modular_pipelines/cosmos/modular_pipeline.py

@@ -79,6 +79,45 @@ def disable_safety_checker(self):
     def requires_safety_checker(self):
         return getattr(self, "_is_safety_checker_enabled", True)
 
+    @property
+    def is_distilled_checkpoint(self) -> bool:
+        """True when the loaded scheduler ships a fixed-step distilled sampler config."""
+        scheduler = getattr(self, "scheduler", None)
+        if scheduler is None:
+            return False
+        fixed_step_cfg = scheduler.config.get("fixed_step_sampler_config", None)
+        requires_explicit_sigmas = bool(scheduler.config.get("fixed_step_requires_explicit_sigmas", False))
+        return bool(fixed_step_cfg is not None and requires_explicit_sigmas and fixed_step_cfg.get("t_list"))
+
+    @property
+    def distilled_sigmas(self) -> list[float]:
+        """Fixed sigma schedule for distilled checkpoints (`fixed_step_sampler_config.t_list`)."""
+        if not self.is_distilled_checkpoint:
+            raise ValueError("distilled_sigmas is only defined for distilled checkpoints.")
+        fixed_step_cfg = self.scheduler.config["fixed_step_sampler_config"]
+        return [float(s) for s in fixed_step_cfg["t_list"]]
+
+    def resolve_inference_schedule(
+        self,
+        num_inference_steps: int | None,
+        guidance_scale: float | None,
+    ) -> tuple[int, float]:
+        """Resolve step count and guidance for base vs distilled checkpoints."""
+        if self.is_distilled_checkpoint:
+            distilled_steps = len(self.distilled_sigmas)
+            if guidance_scale is not None and 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 {guidance_scale})."
+                )
+            return distilled_steps, 1.0
+
+        if guidance_scale is None:
+            guidance_scale = 6.0
+        if num_inference_steps is None:
+            num_inference_steps = 35
+        return num_inference_steps, guidance_scale
+
     def _encode_video(self, x):
         return Cosmos3OmniPipeline._encode_video(self, x)