LoRA fine-tuning for Cosmos Predict 2.5

This example shows how to fine-tune Cosmos Predict 2.5 using LoRA on a custom video dataset.

Requirements

Install the library from source and the example-specific dependencies:

bash

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install -e ".[dev]"
cd examples/cosmos
pip install -r requirements.txt

Data preparation

The training script expects a dataset directory with the following layout:

<dataset_dir>/
├── videos/          # .mp4 files
└── metas/           # one .txt prompt file per video (same stem)
    ├── 0.txt
    ├── 1.txt
    └── ...

GR1 dataset (quick start)

The download_and_preprocess_datasets.sh script downloads the GR1-100 training set and the EVAL-175 test set, then runs the preprocessing script to create the per-video prompt files.

bash

bash download_and_preprocess_datasets.sh

This produces:

gr1_dataset/train/ — training videos + prompts
gr1_dataset/test/ — evaluation images + prompts

Training

Launch LoRA training with accelerate:

bash

export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR="lora-output"

accelerate launch --mixed_precision="bf16" train_cosmos_predict25_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --revision diffusers/base/post-trained \
  --train_data_dir=$DATA_DIR \
  --output_dir=$OUT_DIR \
  --train_batch_size=1 \
  --num_train_epochs=500 \
  --checkpointing_epochs=100 \
  --seed=0 \
  --height 432 --width 768 \
  --allow_tf32 \
  --gradient_checkpointing \
  --lora_rank 32 --lora_alpha 32 \
  --report_to=wandb

Or use the provided shell script:

bash

bash train_lora.sh

Evaluation

Run inference with the trained LoRA adapter:

bash

export DATA_DIR="gr1_dataset/test"
export LORA_DIR="lora-output"
export OUT_DIR="eval-output"

python eval_cosmos_predict25_lora.py \
  --data_dir $DATA_DIR \
  --output_dir $OUT_DIR \
  --lora_dir $LORA_DIR \
  --revision diffusers/base/post-trained \
  --height 432 --width 768 \
  --num_output_frames 93 \
  --num_steps 36 \
  --seed 0

Or use the provided shell script:

bash

bash eval_lora.sh