Merge branch 'main' into autoencoderkl-tests-refactor

Author: sayakpaul

.github/workflows/pr_labeler.yml

@@ -20,6 +20,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
       - name: Check for missing tests
         id: check
         env:

.github/workflows/pr_style_bot.yml

@@ -5,13 +5,14 @@ on:
     types: [created]
 
 permissions:
-  contents: write
   pull-requests: write
+  contents: read
 
 jobs:
   style:
-    uses: huggingface/huggingface_hub/.github/workflows/style-bot-action.yml@e000c1c89c65aee188041723456ac3a479416d4c  # main
+    uses: huggingface/huggingface_hub/.github/workflows/style-bot-action.yml@e2867e92c07d15e1bf18994d0a945ef5ad6b8d65
     with:
       python_quality_dependencies: "[quality]"
     secrets:
-      bot_token: ${{ secrets.HF_STYLE_BOT_ACTION }}
+      app_id: ${{ secrets.HF_BOT_STYLE_APP_ID }}
+      app_private_key: ${{ secrets.HF_BOT_STYLE_SECRET_PEM }}

docs/source/en/api/loaders/lora.md

@@ -132,6 +132,10 @@ LoRA is a fast and lightweight training method that inserts and trains a signifi
 
 [[autodoc]] loaders.lora_pipeline.ZImageLoraLoaderMixin
 
+## CosmosLoraLoaderMixin
+
+[[autodoc]] loaders.lora_pipeline.CosmosLoraLoaderMixin
+
 ## KandinskyLoraLoaderMixin
 [[autodoc]] loaders.lora_pipeline.KandinskyLoraLoaderMixin
 

docs/source/en/optimization/attention_backends.md

@@ -35,7 +35,7 @@ The [`~ModelMixin.set_attention_backend`] method iterates through all the module
 The example below demonstrates how to enable the `_flash_3_hub` implementation for FlashAttention-3 from the [`kernels`](https://github.com/huggingface/kernels) library, which allows you to instantly use optimized compute kernels from the Hub without requiring any setup.
 
 > [!NOTE]
-> FlashAttention-3 is not supported for non-Hopper architectures, in which case, use FlashAttention with `set_attention_backend("flash")`.
+> FlashAttention-3 requires Ampere GPUs at a minimum.
 
 ```py
 import torch

examples/cosmos/README.md

@@ -0,0 +1,97 @@
+# LoRA fine-tuning for Cosmos Predict 2.5
+
+This example shows how to fine-tune [Cosmos Predict 2.5](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B) 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
+```

examples/cosmos/create_prompts_for_gr1_dataset.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from tqdm import tqdm
+
+
+"""example command
+python create_prompts_for_gr1_dataset.py --dataset_path datasets/benchmark_train/gr1
+"""
+
+
+def parse_args() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Create text prompts for GR1 dataset")
+    parser.add_argument(
+        "--dataset_path", type=str, default="datasets/benchmark_train/gr1", help="Root path to the dataset"
+    )
+    parser.add_argument(
+        "--prompt_prefix", type=str, default="The robot arm is performing a task. ", help="Prefix of the prompt"
+    )
+    parser.add_argument(
+        "--meta_csv", type=str, default=None, help="Metadata csv file (defaults to <dataset_path>/metadata.csv)"
+    )
+    return parser.parse_args()
+
+
+def main(args) -> None:
+    meta_csv = args.meta_csv or os.path.join(args.dataset_path, "metadata.csv")
+    meta_lines = open(meta_csv).readlines()[1:]
+    meta_txt_dir = os.path.join(args.dataset_path, "metas")
+    os.makedirs(meta_txt_dir, exist_ok=True)
+
+    for meta_line in tqdm(meta_lines):
+        video_filename, prompt = meta_line.split(",", 1)
+        prompt = prompt.strip("\n")
+        if prompt.startswith('"') and prompt.endswith('"'):
+            # Remove the quotes
+            prompt = prompt[1:-1]
+        prompt = args.prompt_prefix + prompt
+        meta_txt_filename = os.path.join(meta_txt_dir, os.path.basename(video_filename).replace(".mp4", ".txt"))
+        with open(meta_txt_filename, "w") as fp:
+            fp.write(prompt)
+
+        print(f"encoding prompt: {prompt}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/cosmos/download_and_preprocess_datasets.sh

@@ -0,0 +1,25 @@
+dataset_dir='gr1_dataset'
+train_dir=$dataset_dir/train
+test_dir=$dataset_dir/test
+
+# Download and Preprocess Training Dataset
+hf download nvidia/GR1-100 --repo-type dataset --local-dir datasets/benchmark_train/hf_gr1/ && \
+mkdir -p datasets/benchmark_train/gr1/videos && \
+mv datasets/benchmark_train/hf_gr1/gr1/*mp4 datasets/benchmark_train/gr1/videos && \
+mv datasets/benchmark_train/hf_gr1/metadata.csv datasets/benchmark_train/gr1/
+
+python create_prompts_for_gr1_dataset.py --dataset_path datasets/benchmark_train/gr1
+
+# Download Eval Dataset
+hf download nvidia/EVAL-175 --repo-type dataset --local-dir dream_gen_benchmark
+
+
+# Rename dataset directory
+mkdir $dataset_dir
+mv datasets/benchmark_train/gr1 $train_dir
+mv dream_gen_benchmark/gr1_object $test_dir
+echo Download training data to $train_dir
+echo Download test data to $test_dir
+
+# Clean up staging directories
+rm -rf datasets/ dream_gen_benchmark/