Cosmos3 Distilled support by yzhautouskay · Pull Request #14177 · huggingface/diffusers

yzhautouskay · 2026-07-13T15:21:21Z

What does this PR do?

Cosmos3 distilled (few-step) support

Adds inference support for the distilled few-step Cosmos3 checkpoints through the
modular pipeline (Cosmos3OmniModularPipeline).

Checkpoints supported

nvidia/Cosmos3-Super-Text2Image-4Step (distilled T2I)
nvidia/Cosmos3-Super-Image2Video-4Step (distilled I2V)

Both ship a FlowMatchEulerDiscreteScheduler with a fixed sigma schedule
(scheduler.config.fixed_step_sampler_config.t_list) and bake classifier-free
guidance into the weights. They are only supported via the modular pipeline; the
task-based Cosmos3OmniPipeline does not implement the distilled contract.

Before submitting

Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
- Did you read the Coding with AI agents guide?
- Did you self-review the diff against .ai/review-rules.md?
Did you read the contributor guideline?
Did you read our philosophy doc? (important for complex PRs)
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?
Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

yiyixuxu

thanks!
i left some comments

github-actions · 2026-07-14T08:16:10Z

Hi @yzhautouskay, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

yzhautouskay · 2026-07-14T16:54:15Z

thanks @yiyixuxu! I've refactored it to have separate Cosmos3DistilledModularPipeline

yiyixuxu · 2026-07-15T07:21:50Z

@yzhautouskay can you resolve the conflict?

yiyixuxu

thanks, I left some comments

yiyixuxu · 2026-07-15T20:20:56Z

+]
+
+
+def _fake_distilled_components(t_list=(1.0, 0.75, 0.5, 0.25)):


can you update the test to use ModularPipelineTesterMixin?
https://github.com/huggingface/diffusers/blob/main/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py

Done, distilled tests now use ModularPipelineTesterMixin. For now tiny fixture lives under my personal account at yzhautouskay/tiny-cosmos3-distilled-modular-pipe
Could you create hf-internal-testing/tiny-cosmos3-distilled-modular-pipe and copy it over ? I'll then flip the TINY_DISTILLED_REPO constant in the test (and pretrained_model_name_or_path will need to be changed for components in modular_model_index.json)

HuggingFaceDocBuilderDev · 2026-07-16T18:35:19Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

yiyixuxu

thanks, I left some final comments on tests & docs, we can merge after these are addressed

yiyixuxu · 2026-07-16T19:06:47Z

+    return SimpleNamespace(_execution_device="cpu", config=config)
+
+
+def test_cosmos3_distilled_vae_encoder_select_block():


can you move this test to TestCosmos3DistilledModularPipelineFast,
you can do something like

def test_vae_encoder_rejects_image_and_video_together(self): vae_encoder = Cosmos3DistilledBlocks().sub_blocks["vae_encoder"] vae_pipe = vae_encoder.init_pipeline(self.pretrained_model_name_or_path) vae_pipe.load_components(torch_dtype=torch.float32) with pytest.raises(ValueError, match="either image or video"): vae_pipe(image=image, video=video, num_frames=5, height=32, width=32)

yiyixuxu · 2026-07-16T19:13:32Z

+        vae_encoder.select_block(image=object(), video=object())
+
+
+def test_cosmos3_distilled_set_timesteps_declares_distilled_configs():


can you move this test to the TestCosmos3DistilledModularPipelineFast too
Also, it'd be more meaningful to test the final pipeline's config

pipe = self.pipeline_class() assert pipe.config.is_distilled is True assert pipe.config.distilled_sigmas is None

yiyixuxu · 2026-07-16T19:15:11Z

+def test_cosmos3_distilled_text_encoder_requires_str_prompt():
+    Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt="a robot"))
+    with pytest.raises(ValueError, match="`prompt` must be a str"):
+        Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt=["a robot", "another"]))


can you move this test to the TestCosmos3DistilledModularPipelineFast too
also think it is more meaningful to test this behavior on final pipeline, instead of the _check_inputs

def test_rejects_batched_prompts(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs() inputs["prompt"] = ["a robot", "another"] with pytest.raises(ValueError, match="..."): pipe(**inputs, output=self.output_name)

yiyixuxu · 2026-07-16T19:17:33Z

+def test_cosmos3_distilled_text_encoder_omits_negative_prompt():
+    input_names = {inp.name for inp in Cosmos3DistilledTextEncoderStep().inputs}
+    assert "prompt" in input_names
+    assert "negative_prompt" not in input_names


can you remove this test for now? we can extend our base test here to support a "not_params"
https://github.com/huggingface/diffusers/blob/main/tests/modular_pipelines/test_modular_pipelines_common.py#L165

yiyixuxu · 2026-07-16T19:25:35Z

+def _fake_distilled_components(sigmas=(1.0, 0.9375, 0.8333333333333334, 0.625)):
+    config = SimpleNamespace(distilled_sigmas=list(sigmas))
+    return SimpleNamespace(_execution_device="cpu", config=config)


Suggested change

def _fake_distilled_components(sigmas=(1.0, 0.9375, 0.8333333333333334, 0.625)):

config = SimpleNamespace(distilled_sigmas=list(sigmas))

return SimpleNamespace(_execution_device="cpu", config=config)

yiyixuxu · 2026-07-16T19:27:43Z

+        Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt=["a robot", "another"]))
+
+
+def test_cosmos3_distilled_set_timesteps_rejects_step_count_override():


same comment https://github.com/huggingface/diffusers/pull/14177/changes#r3598260664
I think it's more meaningful to test all the behavior on final pipeline (not individual steps)

def test_rejects_num_inference_steps_override(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = 10 with pytest.raises(..."): pipe(**inputs, output=self.output_name)

yiyixuxu · 2026-07-16T19:31:39Z

+        step(_fake_distilled_components(), state)
+
+
+def test_cosmos3_distilled_set_timesteps_rejects_guidance_override():


def test_rejects_guidance_scale_override(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs() inputs["guidance_scale"] = 3.0 with pytest.raises(ValueError, match="`guidance_scale` must be 1.0"): pipe(**inputs, output=self.output_name)

yiyixuxu · 2026-07-16T19:35:58Z

+        step(_fake_distilled_components(), state)
+
+
+def test_cosmos3_distilled_set_timesteps_requires_distilled_sigmas():


maybe remove this one?

yiyixuxu · 2026-07-16T19:36:48Z

+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).


Suggested change

into the weights, `negative_prompt` is not supported (passing one raises an error).

into the weights, `negative_prompt` is not supported (passing one raises a warning and get ignored).

yiyixuxu · 2026-07-16T19:37:11Z

+### 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


I think we need to udpate the doc here, no? now it is a pipeline config

yiyixuxu · 2026-07-16T19:50:11Z

+
+# TODO: move this fixture to `hf-internal-testing/tiny-cosmos3-distilled-modular-pipe` and update the
+# repo name here. Hosted on a personal account for now so the PR can be tested.
+TINY_DISTILLED_REPO = "yzhautouskay/tiny-cosmos3-distilled-modular-pipe"


i created the repo here if you want to PR there https://huggingface.co/hf-internal-testing/tiny-cosmos3-distilled-modular-pipe

github-actions Bot added size/L PR with diff > 200 LOC documentation Improvements or additions to documentation modular-pipelines and removed size/L PR with diff > 200 LOC labels Jul 13, 2026

yiyixuxu reviewed Jul 13, 2026

View reviewed changes

github-actions Bot added tests utils labels Jul 14, 2026

yzhautouskay force-pushed the yzhautouskay/cosmos3_distilled branch from 177ee1e to 3ac7ec0 Compare July 15, 2026 09:43

Cosmos3 distilled support (dedicated modular pipeline + blocks)

abb37e7

yzhautouskay force-pushed the yzhautouskay/cosmos3_distilled branch from 3ac7ec0 to abb37e7 Compare July 15, 2026 09:48

yiyixuxu reviewed Jul 15, 2026

View reviewed changes

yzhautouskay added 2 commits July 16, 2026 03:13

Address review comments

409c2eb

Refactor unit test to use ModularPipelineTesterMixin

6b5199f

yiyixuxu reviewed Jul 16, 2026

View reviewed changes

yiyixuxu mentioned this pull request Jul 16, 2026

Add not_params to ModularPipelineTesterMixin #14207

Open

10 tasks

Merge branch 'main' into yzhautouskay/cosmos3_distilled

235d96a

yiyixuxu mentioned this pull request Jul 16, 2026

[agents doc] notes on when to create new blocksets for checkpoint variant #14208

Open

		]


		def _fake_distilled_components(t_list=(1.0, 0.75, 0.5, 0.25)):

		return SimpleNamespace(_execution_device="cpu", config=config)


		def test_cosmos3_distilled_vae_encoder_select_block():

		vae_encoder.select_block(image=object(), video=object())


		def test_cosmos3_distilled_set_timesteps_declares_distilled_configs():

	def _fake_distilled_components(sigmas=(1.0, 0.9375, 0.8333333333333334, 0.625)):
	config = SimpleNamespace(distilled_sigmas=list(sigmas))
	return SimpleNamespace(_execution_device="cpu", config=config)

		Cosmos3DistilledTextEncoderStep._check_inputs(SimpleNamespace(prompt=["a robot", "another"]))


		def test_cosmos3_distilled_set_timesteps_rejects_step_count_override():

		step(_fake_distilled_components(), state)


		def test_cosmos3_distilled_set_timesteps_rejects_guidance_override():

	into the weights, `negative_prompt` is not supported (passing one raises an error).
	into the weights, `negative_prompt` is not supported (passing one raises a warning and get ignored).

Uh oh!

Conversation

yzhautouskay commented Jul 13, 2026

What does this PR do?

Cosmos3 distilled (few-step) support

Checkpoints supported

Before submitting

Who can review?

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jul 14, 2026

Uh oh!

yzhautouskay commented Jul 14, 2026

Uh oh!

yiyixuxu commented Jul 15, 2026

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Jul 16, 2026

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants