Skip to content

Add transfer support to the Cosmos3 modular pipeline#1

Closed
yzhautouskay wants to merge 25 commits into
yzhautouskay/cosmos3_modular_pipelinefrom
yzhautouskay/cosmos3_modular_transfer
Closed

Add transfer support to the Cosmos3 modular pipeline#1
yzhautouskay wants to merge 25 commits into
yzhautouskay/cosmos3_modular_pipelinefrom
yzhautouskay/cosmos3_modular_transfer

Conversation

@yzhautouskay

@yzhautouskay yzhautouskay commented Jul 8, 2026

Copy link
Copy Markdown
Owner

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.

How it works

The pipeline accepts a control_videos mapping ({hint: video}); each control clip is VAE-encoded and injected as an additional clean conditioning item before the noisy target during denoising, so generated frames follow the control's spatial structure while the prompt drives appearance. A separate control_guidance (and optional control_guidance_interval) scales how strongly the structural signal is enforced, independently of the text guidance_scale, via a nested control/text CFG blend. It runs in V2V mode with num_conditional_frames priming frames, generating long clips autoregressively in chunks of num_video_frames_per_chunk and stitching them into a temporally coherent multi-second clip

Supported hints

edge, blur, depth, seg, wsm

Notes for reviewers

This PR is stacked on the Cosmos3 modular pipeline PR 14110 the diff here is transfer-only relative to that branch. It will be rebased onto main once the modular PR merges.

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.

dependabot Bot and others added 20 commits July 2, 2026 17:36
…uggingface#14109)

Bumps [transformers](https://github.com/huggingface/transformers) from 4.47.0 to 5.3.0.
- [Release notes](https://github.com/huggingface/transformers/releases)
- [Commits](huggingface/transformers@v4.47.0...v5.3.0)

---
updated-dependencies:
- dependency-name: transformers
  dependency-version: 5.3.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Fix mutable default args in LoRA mixin

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
…ace#14096)

* Support root-hosted weights + add _allow_files download hook

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
)

* fix test_from_save_pretrained_dtype_inference

* fix

* add warning to to()

* address feedback

* fix

* fix

* fix

* logger.debug -> logger.warning

* Update src/diffusers/models/modeling_utils.py

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

* fix

---------

Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
…uggingface#14091)

Fix AutoencoderVidTok output format and skip training tests on unsupported hardware
* fix fa3 mask

* add test case

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Update doc build workflow: light installs, drop custom container

* Update doc build workflow: light installs, drop custom container

* Keep the pin comment dependabot-compatible

* Keep the pin comment dependabot-compatible
* fix tests for transformers latest stable.

* fix

---------

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
fix rf-time scheduile problem

Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* feat: add image edit plus

* refactor: remove debug code

* fix: address review issues for JoyImage Edit Plus

    - Remove einops dependency: replace rearrange with reshape/permute
    - Remove sglang-specific code from transformer forward
    - Remove unused import inspect from transformer
    - Fix hardcoded device_type="cuda" to use device.type
    - Simplify scheduler sigma math: delegate to retrieve_timesteps
    - Remove unused enable_denormalization parameter
    - Fix callback latents variable binding
    - Fix output_type="pt" to return stacked tensor
    - Set return_dict default to True in transformer forward
    - Add dummy objects for JoyImageEditPlus classes
    - Add transformer and pipeline test files

* fix: add missing newline at end of pipeline_output.py

* fix: add missing newline at end of pipeline_output.py

* doc: add joyimage-edit-plus doc

* refactor: update code format

* refactor: update code formate

* refactor: merge edit-plus conversion script into convert_joyimage_edit_to_diffusers.py

	JoyImage Edit and Edit Plus share identical VAE and transformer weight
      layouts — only the target model class differs. Consolidate both into a
      single script with a --model_type flag (edit | edit_plus) instead of
      maintaining two nearly-duplicate files.

* revert: remove unnecessary attention_mask change from transformer_joyimage.py

* Apply style fixes

* fix: remove stale Copied-from annotations that diverged from source

* fix: address CI check failures for edit-plus PR

  - Add missing Returns section to JoyImageEditPlusTransformer3DModel.forward docstring
  - Fix alphabetical ordering of dummy classes in dummy_pt_objects.py

* fix: use VAE dtype instead of float32 cast in edit-plus pipeline

   Replace `torch.autocast(dtype=torch.float32)` + `.float()` with
   `.to(self.vae.dtype)` for both VAE encode and decode calls.

   The previous approach caused dtype mismatch (float32 input vs bfloat16
   bias) on CPU where autocast does not automatically cast conv weights,
   breaking the CI `test_layerwise_casting_inference` test.

* refactor: address dg845 review comments on edit-plus

  Transformer:
  - Simplify RoPE by removing dead txt_freqs path
  - Fix patch_size type hint to list[int]
  - Use self.config.xxx instead of explicit attributes

  Pipeline:
  - Merge VaeImageProcessor into JoyImageEditImageProcessor
  - Inline normalize/denormalize latents at call sites
  - Use image_processor.preprocess() for reference images
  - Remove duplicate resize_center_crop in prepare_latents
  - Rename padded_latents to latents in __call__
  - Support pre-computed latents parameter
  - Remove redundant padding token zeroing (attention mask suffices)
  - Use image_processor.postprocess() for output conversion
  - Return (image,) tuple when return_dict=False

* fix: guard against division-by-zero in CFG norm rescaling

  Add .clamp_min(1e-6) to noise_norm to prevent NaN when comb_pred
  is all-zeros (e.g. padded patches), matching pipeline_joyimage_edit.py.

* style: run ruff format on edit-plus pipeline

* refactor: address dg845 second-round review comments

  - Fix vae_image_processor -> image_processor in docs
  - Remove empty __init__ from AttnProcessor
  - Restore # Copied from on retrieve_timesteps and sync with source
  - Add check_inputs method for basic input validation
  - Remove unnecessary test skip for test_from_save_pretrained_dtype_inference

* Apply style fixes

---------

Co-authored-by: tangyanfei.8 <tangyanfei.8@jd.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
fix

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* add `_flash_3_varlen_hub` support for SP

* make signiture clear

* update

* add accuracy test

* fix dtype

* loosen threshold because of bfloat16

* fix the break change of huggingface/kernels-community@1471f94

* fix stale error message

* fix merge

* revert change

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* post release v0.39.0

* style.

* remove deprecated code path.

* remove deprecated method.
This reverts commit b3545f9.
@github-actions github-actions Bot added documentation Improvements or additions to documentation modular-pipelines labels Jul 8, 2026
@yzhautouskay yzhautouskay force-pushed the yzhautouskay/cosmos3_modular_transfer branch from f8ba319 to 3563da1 Compare July 9, 2026 10:58
yzhautouskay and others added 4 commits July 9, 2026 08:26
* Cosmos3 ModularPipeline initial commit

* Fix from_pretrained for modular pipleine without modular_model_index.json

* Isolate text prep inside Cosmos3TextEncoderStep; add Cosmos3ActionTextStep

* Cosmos3 modular encode/decode flow

* Drop device/dtype required InputParams; use components._execution_device and model dtype

* Add lazy pipeline properties and requires_safety_checker-based safety gating

* Drop callback_on_step_end from Cosmos3 denoise loop

* Move container/auto blocks into modular_blocks

* Move action_gen check, remove unused params

* Declare video_processor as ComponentSpec; drop unused pipeline state

* Split VAE encode block by modality (image and video)

* Remove Cosmos3AssembleResultStep; add Cosmos3ExportStep

* Remove unused action_mode

* Refactor Cosmos3 modular denoising by modality

* Document Cosmos3 modular output selection

* Address Cosmos3 modular review feedback

* Remaining local-variable device cleanup

* Address Cosmos3 modular review feedback

* Address final Cosmos3 modular feedback

---------

Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
@yzhautouskay yzhautouskay force-pushed the yzhautouskay/cosmos3_modular_transfer branch from 27783f0 to 2181295 Compare July 9, 2026 19:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.