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: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: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: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

Thanks, opened PR: https://huggingface.co/hf-internal-testing/tiny-cosmos3-distilled-modular-pipe/discussions/1

yiyixuxu

thanks

… checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields).

Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

… checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields). Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

… checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields). Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

… checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields). Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

… checkpoint Register nvidia/Cosmos3-Super-Image2Video-4Step and add the one algorithmic piece distilled I2V needs: the stochastic FlowMatchEuler step re-noises every position each step, so the clean conditioning frame is re-anchored after every scheduler step via the denoise post_step_fn hook (matching the diffusers distilled loop, PR huggingface/diffusers#14177). Base UniPC sampling is unchanged: deterministic steps never move a zero-velocity frame. This lifts the temporary rejection of image-conditioned requests on distilled checkpoints, which was added to hold the line until per-step re-anchoring landed. Also raise on enable_audio=True when the checkpoint ships no audio tower (weight-presence guard, not workflow policy), make the final conditioning re-injection in-place instead of cloning the full latent tensor, and accept the I2V-4Step transformer config shape (sound_dim null, no action fields). Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

Smoke: run the documented example invocation (deterministic PIL conditioning image, omni-default 720p x 189 shape) and assert a non-empty MP4. Quality gate: unlike the existing TRT-LLM self-goldens, the golden video is produced by the reference implementation (diffusers Cosmos3 distilled modular pipeline, huggingface/diffusers#14177, with its per-step SDE noise made generator-seeded), so the gate checks the denoising trajectory against the reference rather than regression against a past TRT-LLM run. Full provenance (diffusers commit, RNG patch, corrected modular index, generation parameters) is recorded in cosmos3_i2v_4step_lpips_golden_video.json. Threshold 0.10 = 0.0563 measured at golden creation plus headroom for the ~0.04 cross-host kernel drift documented in the harness; validated at 0.0588 on B200. Signed-off-by: Igor Shovkun <ishovkun@nvidia.com>

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

Merged

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

Merged

yzhautouskay added 3 commits July 17, 2026 02:35

Docs fix

18ca731

Tests refactor

c09807a

Pass generator to SDE scheduler for reproducibility

874a05e

yiyixuxu reviewed Jul 17, 2026

View reviewed changes

Comment thread tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py Outdated

yiyixuxu added 2 commits July 17, 2026 06:52

Apply suggestion from @yiyixuxu

12ac6c2

Merge branch 'main' into yzhautouskay/cosmos3_distilled

88d7f3d

yiyixuxu approved these changes Jul 17, 2026

View reviewed changes

Distilled text encoder step bugfix

1ea951c

yiyixuxu merged commit 8eac6ad into huggingface:main Jul 17, 2026
19 checks passed

ishovkun mentioned this pull request Jul 21, 2026

[None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint NVIDIA/TensorRT-LLM#16690

Merged

1 task

		]


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

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


		def test_cosmos3_distilled_set_timesteps_declares_distilled_configs():

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


		def test_cosmos3_distilled_set_timesteps_rejects_step_count_override():

Uh oh!

Conversation

yzhautouskay commented Jul 13, 2026 • edited by yiyixuxu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

yiyixuxu Jul 15, 2026

Choose a reason for hiding this comment

Uh oh!

yzhautouskay Jul 16, 2026

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!

Uh oh!

yiyixuxu Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiyixuxu Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

yzhautouskay Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yzhautouskay commented Jul 13, 2026 •

edited by yiyixuxu

Loading