Add LoRA support for Cosmos Predict 2.5 and fix pipeline to match official Cosmos repo by terarachang · Pull Request #13664 · huggingface/diffusers

terarachang · 2026-04-30T20:12:17Z

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)

…ncoder attention implementation, and timestep scaling

… to device before torch.stack

yiyixuxu

thanks, i left a question

yiyixuxu · 2026-05-05T23:03:26Z


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


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

yiyixuxu

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

sayakpaul

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

HuggingFaceDocBuilderDev · 2026-05-07T20:49:33Z

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.

…icial Cosmos repo (huggingface#13664) * support lora for cosmos 2.5 * Fix inconsistencies with cosmos official repo in VAE encoding, text encoder attention implementation, and timestep scaling * Support f_min and f_max in linear_scheduler warmup * Add requirements and dataset preprocessing scripts to run examples * Add LoRA training scripts * Add LoRA eval scripts * add assets for blogpost * Fix(scheduler): device mismatch from upstream b114620 - move rk and b to device before torch.stack * Always upcast to fp32 * Directly inhrit from LoraBaseMixin * remove flash-attn2 * Use _keep_in_fp32_modules instead of autocast * remove the get_latent_shape_cthw method and fix style * simplifiy the eval script to make it more user-friendly * overwrite scheduling_unipc_multistep.py with main's version * remove network_alphas and add # Copied from * remove figures and assets * revert scheduler * revert fp32 upcast and support bs > 1 --------- Co-authored-by: Ting-Yun Chang <tingyunc@nvidia.com>

Ting-Yun Chang added 8 commits April 30, 2026 19:52

support lora for cosmos 2.5
48fbe84

Fix inconsistencies with cosmos official repo in VAE encoding, text e…
5216816
…ncoder attention implementation, and timestep scaling

Support f_min and f_max in linear_scheduler warmup
3d5551b

Add requirements and dataset preprocessing scripts to run examples
f0b51ed

Add LoRA training scripts
9038e10

Add LoRA eval scripts
cee46cc

add assets for blogpost
9e24319

Fix(scheduler): device mismatch from upstreamb114620- move rk and b…
c3bb1aa
… to device before torch.stack

github-actionsBot added documentation Improvements or additions to documentation lora models pipelines examples schedulers loaders size/L PR with diff > 200 LOC labels Apr 30, 2026

Always upcast to fp32
d32b085

github-actionsBot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 30, 2026

Directly inhrit from LoraBaseMixin
80aa547

github-actionsBot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels May 2, 2026