Skip to content

Add LoRA support for Cosmos Predict 2.5 and fix pipeline to match official Cosmos repo#13664

Merged
yiyixuxu merged 19 commits intohuggingface:mainfrom
terarachang:cosmos_predict_2.5_lora_clean
May 8, 2026
Merged

Add LoRA support for Cosmos Predict 2.5 and fix pipeline to match official Cosmos repo#13664
yiyixuxu merged 19 commits intohuggingface:mainfrom
terarachang:cosmos_predict_2.5_lora_clean

Conversation

@terarachang
Copy link
Copy Markdown
Contributor

@terarachang terarachang commented Apr 30, 2026

What this PR does

Adds LoRA fine-tuning support for Cosmos Predict 2.5 (nvidia/Cosmos-Predict2.5-2B) and fixes the pipeline to match the official Cosmos reference implementation.

LoRA support

  • CosmosLoraLoaderMixin in src/diffusers/loaders/lora_pipeline.py for LoRA loading/saving on CosmosTransformer3DModel
  • Training script examples/cosmos/train_cosmos_predict25_lora.py using accelerate + peft
  • Inference script examples/cosmos/eval_cosmos_predict25_lora.py
  • Added CosmosLoraLoaderMixin to docs/source/en/api/loaders/lora.md

Fixes to match the official Cosmos repo

  • Fix conditional_frame_timestep scaling by timestep_scale=0.001
  • Auto-cast AdaLN and DiT final layer to fp32 for training
  • Deterministic VAE encode (no sampling)

Test plan

  • All repository consistency checks pass (check_copies, check_dummies, check_support_list)

@github-actions github-actions Bot added documentation Improvements or additions to documentation lora models pipelines examples schedulers loaders size/L PR with diff > 200 LOC labels Apr 30, 2026
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 30, 2026
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 2, 2026
Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

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

thanks, i left a question

Comment thread src/diffusers/models/transformers/transformer_cosmos.py Outdated
Comment thread src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py Outdated
Comment thread src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py Outdated
Comment thread src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py Outdated
Comment thread src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py Outdated
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 5, 2026
@terarachang terarachang force-pushed the cosmos_predict_2.5_lora_clean branch from 0cc6351 to c8513c2 Compare May 5, 2026 19:25
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 5, 2026
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 5, 2026

device = sample.device
sigma_t, sigma_s0 = self.sigmas[self.step_index + 1].to(device), self.sigmas[self.step_index].to(device)
sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
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.

ohh I think the change here might not be intended, no?
it seeem to have reverted https://github.com/huggingface/diffusers/pull/13489/changes

Copy link
Copy Markdown
Collaborator

@yiyixuxu yiyixuxu left a comment

Choose a reason for hiding this comment

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

thanks, looks good to me once we revert the change in scheculer

@yiyixuxu yiyixuxu requested a review from sayakpaul May 5, 2026 23:06
@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 5, 2026
Copy link
Copy Markdown
Member

@sayakpaul sayakpaul left a comment

Choose a reason for hiding this comment

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

I left some questions and suggestions. LMK if anything is unclear.

Comment thread docs/source/en/api/loaders/lora.md
Comment thread src/diffusers/loaders/lora_pipeline.py
Comment thread src/diffusers/loaders/lora_pipeline.py
Comment thread src/diffusers/loaders/lora_pipeline.py
Comment thread src/diffusers/loaders/lora_pipeline.py Outdated
Comment thread src/diffusers/loaders/lora_pipeline.py
@HuggingFaceDocBuilderDev
Copy link
Copy Markdown

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 yiyixuxu merged commit 4ca8633 into huggingface:main May 8, 2026
34 of 36 checks passed
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 examples loaders lora models pipelines size/L PR with diff > 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants