Skip to content

Cosmos3 Distilled support#14177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yiyixuxu merged 10 commits into from
Jul 17, 2026
Merged

Cosmos3 Distilled support #14177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
abb37e7
Cosmos3 distilled support (dedicated modular pipeline + blocks)
yzhautouskay Jul 15, 2026
409c2eb
Address review comments
yzhautouskay Jul 16, 2026
6b5199f
Refactor unit test to use ModularPipelineTesterMixin
yzhautouskay Jul 16, 2026
235d96a
Merge branch 'main' into yzhautouskay/cosmos3_distilled
yiyixuxu Jul 16, 2026
18ca731
Docs fix
yzhautouskay Jul 17, 2026
c09807a
Tests refactor
yzhautouskay Jul 17, 2026
874a05e
Pass generator to SDE scheduler for reproducibility
yzhautouskay Jul 17, 2026
12ac6c2
Apply suggestion from @yiyixuxu
yiyixuxu Jul 17, 2026
88d7f3d
Merge branch 'main' into yzhautouskay/cosmos3_distilled
yiyixuxu Jul 17, 2026
1ea951c
Distilled text encoder step bugfix
yzhautouskay Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions docs/source/en/api/pipelines/cosmos3.md
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down
4 changes: 4 additions & 0 deletions src/diffusers/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -489,6 +489,8 @@
[
"AnimaAutoBlocks",
"AnimaModularPipeline",
"Cosmos3DistilledBlocks",
"Cosmos3DistilledModularPipeline",
"Cosmos3OmniBlocks",
"Cosmos3OmniModularPipeline",
"ErnieImageAutoBlocks",
Expand DownExpand Up@@ -1349,6 +1351,8 @@
from .modular_pipelines import (
AnimaAutoBlocks,
AnimaModularPipeline,
Cosmos3DistilledBlocks,
Cosmos3DistilledModularPipeline,
Cosmos3OmniBlocks,
Cosmos3OmniModularPipeline,
ErnieImageAutoBlocks,
Expand Down
9 changes: 8 additions & 1 deletion src/diffusers/modular_pipelines/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -98,6 +98,8 @@
"AnimaModularPipeline",
]
_import_structure["cosmos"] = [
"Cosmos3DistilledBlocks",
"Cosmos3DistilledModularPipeline",
"Cosmos3OmniBlocks",
"Cosmos3OmniModularPipeline",
]
Expand DownExpand Up@@ -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 (
Expand Down
6 changes: 4 additions & 2 deletions src/diffusers/modular_pipelines/cosmos/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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:
Expand All@@ -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

Expand Down
94 changes: 93 additions & 1 deletion src/diffusers/modular_pipelines/cosmos/before_denoise.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand DownExpand Up@@ -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()
Expand DownExpand Up@@ -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
Expand DownExpand Up@@ -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),
Comment thread
yzhautouskay marked this conversation as resolved.
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
88 changes: 87 additions & 1 deletion src/diffusers/modular_pipelines/cosmos/denoise.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand DownExpand Up@@ -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"

Expand DownExpand Up@@ -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,
Expand Down
Loading
Loading