Cosmos3 distilled refactor for separate modualr pipeline

Author: yzhautouskay

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

@@ -938,13 +938,14 @@ if outputs["action"] is not None:
 
 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.
+the weights. They are served by the dedicated [`Cosmos3DistilledModularPipeline`] (blocks:
+`Cosmos3DistilledBlocks`) — the task-based [`Cosmos3OmniPipeline`] and the base
+[`Cosmos3OmniModularPipeline`] do 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.
+Load a distilled repo and call the pipeline **without** `num_inference_steps` or
+`guidance_scale`; `Cosmos3DistilledSetTimestepsStep` reads the fixed step count from the
+scheduler config and forces `guidance_scale=1.0`. Because classifier-free guidance is baked
+into the weights, `negative_prompt` is not supported (passing one raises an error).
 
 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
@@ -953,14 +954,14 @@ described in [Prompt upsampling](#prompt-upsampling), then pass the JSON via `js
 ```python
 import json
 import torch
-from diffusers import Cosmos3OmniModularPipeline
+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 = Cosmos3OmniModularPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
+pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16)
 pipe.load_components(torch_dtype=torch.bfloat16)
 pipe.to("cuda")
 
@@ -980,7 +981,7 @@ videos[0].save("cosmos3_distilled_t2i.jpg", format="JPEG", quality=85)
 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 = Cosmos3DistilledModularPipeline.from_pretrained(repo_i2v, torch_dtype=torch.bfloat16)
 pipe.load_components(torch_dtype=torch.bfloat16)
 pipe.to("cuda")
 
@@ -1006,6 +1007,13 @@ For a CLI wrapper with warmup / iteration timing, see
 - all
 - __call__
 
+## Cosmos3DistilledModularPipeline
+
+[[autodoc]] Cosmos3DistilledModularPipeline
+
+- all
+- __call__
+
 ## CosmosActionCondition
 
 [[autodoc]] CosmosActionCondition

src/diffusers/__init__.py

@@ -477,6 +477,8 @@
         [
             "AnimaAutoBlocks",
             "AnimaModularPipeline",
+            "Cosmos3DistilledBlocks",
+            "Cosmos3DistilledModularPipeline",
             "Cosmos3OmniBlocks",
             "Cosmos3OmniModularPipeline",
             "ErnieImageAutoBlocks",
@@ -1329,6 +1331,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

@@ -4,10 +4,10 @@
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
 from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
-from ...schedulers import SchedulerMixin
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ...utils.torch_utils import randn_tensor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
-from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
@@ -187,7 +187,7 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("transformer", Cosmos3OmniTransformer),
-            ComponentSpec("scheduler", SchedulerMixin),
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
     @property
@@ -263,7 +263,7 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [
             ComponentSpec("transformer", Cosmos3OmniTransformer),
-            ComponentSpec("scheduler", SchedulerMixin),
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
         ]
 
     @property
@@ -940,7 +940,48 @@ def description(self) -> str:
 
     @property
     def expected_components(self) -> list[ComponentSpec]:
-        return [ComponentSpec("scheduler", SchedulerMixin)]
+        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_inference_steps", required=True),
+        ]
+
+    @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."),
+        ]
+
+    @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
+        )
+        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 scheduler's `fixed_step_sampler_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)]
 
     @property
     def inputs(self) -> list[InputParam]:
@@ -950,7 +991,10 @@ def inputs(self) -> list[InputParam]:
                 name="guidance_scale",
                 type_hint=float,
                 default=None,
-                description="Classifier-free guidance scale. Forced to 1.0 for distilled checkpoints.",
+                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."
+                ),
             ),
         ]
 
@@ -962,32 +1006,46 @@ def intermediate_outputs(self) -> list[OutputParam]:
             OutputParam(
                 "num_inference_steps",
                 type_hint=int,
-                description="Resolved number of denoising steps.",
+                description="Resolved number of denoising steps (fixed by the distilled schedule).",
             ),
             OutputParam(
                 name="guidance_scale",
                 type_hint=float,
-                description="Resolved classifier-free guidance scale for the denoising loop.",
+                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
-        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)
+        fixed_step_cfg = components.scheduler.config.get("fixed_step_sampler_config", None)
+        if not fixed_step_cfg or not fixed_step_cfg.get("t_list"):
+            raise ValueError(
+                "Cosmos3DistilledSetTimestepsStep requires a scheduler that ships a distilled "
+                "`fixed_step_sampler_config.t_list`. Load a distilled Cosmos3 checkpoint or use "
+                "`Cosmos3OmniModularPipeline` for base checkpoints."
+            )
+        sigmas = [float(s) for s in fixed_step_cfg["t_list"]]
+        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 scheduler's "
+                f"`fixed_step_sampler_config.t_list` ({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) - num_inference_steps * components.scheduler.order
+        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 SchedulerMixin
+from ...schedulers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler
 from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 from .modular_pipeline import Cosmos3OmniModularPipeline
@@ -254,7 +254,39 @@ def description(self) -> str:
 
     @property
     def expected_components(self) -> list[ComponentSpec]:
-        return [ComponentSpec("scheduler", SchedulerMixin)]
+        return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
+
+    @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."
+            ),
+        ]
+
+    @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):
+        block_state.latents = components.scheduler.step(
+            block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False
+        )[0].squeeze(0)
+        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]:
@@ -296,11 +328,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
         # 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
-        ):
+        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
@@ -326,7 +354,7 @@ def inputs(self) -> list[InputParam]:
             ),
             InputParam(
                 name="sound_scheduler",
-                type_hint=SchedulerMixin,
+                type_hint=UniPCMultistepScheduler,
                 required=True,
                 description="Scheduler used to update sound latents.",
             ),
@@ -365,7 +393,7 @@ def inputs(self) -> list[InputParam]:
             ),
             InputParam(
                 name="action_scheduler",
-                type_hint=SchedulerMixin,
+                type_hint=UniPCMultistepScheduler,
                 required=True,
                 description="Scheduler used to update action latents.",
             ),
@@ -412,7 +440,7 @@ def description(self) -> str:
     @property
     def loop_expected_components(self) -> list[ComponentSpec]:
         return [
-            ComponentSpec("scheduler", SchedulerMixin),
+            ComponentSpec("scheduler", UniPCMultistepScheduler),
             ComponentSpec("transformer", Cosmos3OmniTransformer),
         ]
 
@@ -453,6 +481,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,