Skip to content

Transfer support for Cosmos3 ModularPipeline#14150

Open
yzhautouskay wants to merge 5 commits into
huggingface:mainfrom
yzhautouskay:yzhautouskay/cosmos3_modular_transfer
Open

Transfer support for Cosmos3 ModularPipeline#14150
yzhautouskay wants to merge 5 commits into
huggingface:mainfrom
yzhautouskay:yzhautouskay/cosmos3_modular_transfer

Conversation

@yzhautouskay

@yzhautouskay yzhautouskay commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds the transfer (structure-controlled) path to the Cosmos3 Modular pipeline (Cosmos3OmniModularPipeline / Cosmos3OmniBlocks), letting users condition video generation on one or more control videos (segmentation, depth, edge/canny, blur, world-scenario map) in addition to the text prompt. This mirrors the transfer path in the task-based Cosmos3OmniPipeline and is element-wise consistent with it

What this adds

  • A transfer workflow for the modular pipeline, triggered automatically when control_videos is passed -- no new pipeline class, it coexists with the existing modular modes (text2image, text2video, image2video, video2video, action, sound).
  • Transfer is implemented as first-class modular blocks (matching how the existing modalities are built), not a monolithic step:
    • text: Cosmos3TransferTextStep
    • before-denoise: Cosmos3TransferSetupStep, Cosmos3TransferPrepareLatentsStep, Cosmos3TransferPackSequenceStep, Cosmos3TransferSetTimestepsStep
    • denoise: per-chunk loop blocks (Cosmos3TransferChunkVaeEncoderStep -> prepare -> denoise -> scheduler), with the denoiser consuming the packed static contexts through the shared denoiser_input_fields mechanism
    • decode: Cosmos3TransferDecodeChunkStep, Cosmos3TransferStitchStep
    • composed/wired in modular_blocks_cosmos3.py via the Auto* blocks (selected on control_videos), with "transfer" added to available_workflows.

Before submitting

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.

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

@yiyixuxu yiyixuxu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks a lot, i left some feedbacks

Comment thread src/diffusers/modular_pipelines/cosmos/encoders.py Outdated
Comment thread src/diffusers/modular_pipelines/cosmos/before_denoise.py Outdated
]

@torch.no_grad()
def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you consider using the same pattern as Helio for its chunk loop? https://github.com/huggingface/diffusers/blob/main/src/diffusers/modular_pipelines/helios/denoise.py

Basically, we make the outer chunk loop a proper LoopSequentialPipelineBlocks (see https://github.com/huggingface/diffusers/blob/main/src/diffusers/modular_pipelines/helios/denoise.py#L706) and write the inner denoiser loop as a plain block (see https://github.com/huggingface/diffusers/blob/main/src/diffusers/modular_pipelines/helios/denoise.py#L355)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I have looked at Helios pipeline, and the constraint was that LoopSequentialPipelineBlocks rejects non-leaf sub-blocks, and our inner timestep loop (Cosmos3TransferDenoiseStep -> Cosmos3DenoiseLoopWrapper) is itself a LoopSequentialPipelineBlocks, so the two loops can't both be loop-wrappers

And I went with keeping the inner loop as the decomposed LoopSequentialPipelineBlocks because that's the structure every other cosmos modality uses (vision/sound/action all subclass Cosmos3DenoiseLoopWrapper), leaving the chunk loop as the only transfer-specific piece, doing it with __call__ override of SequentialPipelineBlocks

Your proposal also works, but the transfer's inner denoise becomes a monolithic leaf block, diverging from the other modalities (cross-pipeline consistency with Helios vs. intra-cosmos consistency)

Let me know if I am missing something, and you would prefer the refactor

@github-actions

Copy link
Copy Markdown
Contributor

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.

@yiyixuxu yiyixuxu added the no-issue-needed for PRs that do not require link to an issue label Jul 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation modular-pipelines no-issue-needed for PRs that do not require link to an issue size/L PR with diff > 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants