update

.github/workflows/build_docker_images.yml

@@ -38,16 +38,16 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Build Changed Docker Images
-        env: 
-          CHANGED_FILES: ${{ steps.file_changes.outputs.all }}
+        env:
+          CHANGED_FILES:  "${{ steps.file_changes.outputs.all }}"
         run: |
-          echo "$CHANGED_FILES"
-          for FILE in $CHANGED_FILES; do 
-            # skip anything that isn't still on disk
+          for FILE in $CHANGED_FILES; do
+            # skip anything that isn’t still on disk
             if [[ ! -f "$FILE" ]]; then
               echo "Skipping removed file $FILE"
               continue
-            fi           
+            fi
+            
             if [[ "$FILE" == docker/*Dockerfile ]]; then
               DOCKER_PATH="${FILE%/Dockerfile}"
               DOCKER_TAG=$(basename "$DOCKER_PATH")

.github/workflows/nightly_tests.yml

@@ -13,9 +13,8 @@ env:
   PYTEST_TIMEOUT: 600
   RUN_SLOW: yes
   RUN_NIGHTLY: yes
-  PIPELINE_USAGE_CUTOFF: 0
+  PIPELINE_USAGE_CUTOFF: 5000
   SLACK_API_TOKEN: ${{ secrets.SLACK_CIFEEDBACK_BOT_TOKEN }}
-  CONSOLIDATED_REPORT_PATH: consolidated_test_report.md
 
 jobs:
   setup_torch_cuda_pipeline_matrix:
@@ -100,6 +99,11 @@ jobs:
         with:
           name: pipeline_${{ matrix.module }}_test_reports
           path: reports
+      - name: Generate Report and Notify Channel
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
 
   run_nightly_tests_for_other_torch_modules:
     name: Nightly Torch CUDA Tests
@@ -138,6 +142,7 @@ jobs:
         HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }}
         # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
         CUBLAS_WORKSPACE_CONFIG: :16:8
+        RUN_COMPILE: yes
       run: |
         python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
           -s -v -k "not Flax and not Onnx" \
@@ -170,6 +175,12 @@ jobs:
         name: torch_${{ matrix.module }}_cuda_test_reports
         path: reports
 
+    - name: Generate Report and Notify Channel
+      if: always()
+      run: |
+        pip install slack_sdk tabulate
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
   run_torch_compile_tests:
     name: PyTorch Compile CUDA tests
 
@@ -213,6 +224,12 @@ jobs:
         name: torch_compile_test_reports
         path: reports
 
+    - name: Generate Report and Notify Channel
+      if: always()
+      run: |
+        pip install slack_sdk tabulate
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+  
   run_big_gpu_torch_tests:
     name: Torch tests on big GPU
     strategy:
@@ -263,7 +280,12 @@ jobs:
         with:
           name: torch_cuda_big_gpu_test_reports
           path: reports
-
+      - name: Generate Report and Notify Channel
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+          
   torch_minimum_version_cuda_tests:
     name: Torch Minimum Version CUDA Tests
     runs-on:
@@ -320,6 +342,63 @@ jobs:
         with:
           name: torch_minimum_version_cuda_test_reports
           path: reports
+ 
+  run_flax_tpu_tests:
+    name: Nightly Flax TPU Tests
+    runs-on:
+      group: gcp-ct5lp-hightpu-8t
+    if: github.event_name == 'schedule'
+
+    container:
+      image: diffusers/diffusers-flax-tpu
+      options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} -v /mnt/hf_cache:/mnt/hf_cache
+    defaults:
+      run:
+        shell: bash
+    steps:
+    - name: Checkout diffusers
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 2
+
+    - name: Install dependencies
+      run: |
+        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
+        python -m uv pip install -e [quality,test]
+        pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
+        python -m uv pip install pytest-reportlog
+
+    - name: Environment
+      run: python utils/print_env.py
+
+    - name: Run nightly Flax TPU tests
+      env:
+        HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }}
+      run: |
+        python -m pytest -n 0 \
+          -s -v -k "Flax" \
+          --make-reports=tests_flax_tpu \
+          --report-log=tests_flax_tpu.log \
+          tests/
+
+    - name: Failure short reports
+      if: ${{ failure() }}
+      run: |
+        cat reports/tests_flax_tpu_stats.txt
+        cat reports/tests_flax_tpu_failures_short.txt
+
+    - name: Test suite reports artifacts
+      if: ${{ always() }}
+      uses: actions/upload-artifact@v4
+      with:
+        name: flax_tpu_test_reports
+        path: reports
+
+    - name: Generate Report and Notify Channel
+      if: always()
+      run: |
+        pip install slack_sdk tabulate
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
 
   run_nightly_onnx_tests:
     name: Nightly ONNXRuntime CUDA tests on Ubuntu
@@ -370,12 +449,18 @@ jobs:
         name: tests_onnx_cuda_reports
         path: reports
 
+    - name: Generate Report and Notify Channel
+      if: always()
+      run: |
+        pip install slack_sdk tabulate
+        python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
   run_nightly_quantization_tests:
     name: Torch quantization nightly tests
     strategy:
       fail-fast: false
       max-parallel: 2
-      matrix:
+      matrix: 
         config:
           - backend: "bitsandbytes"
             test_location: "bnb"
@@ -435,7 +520,12 @@ jobs:
         with:
           name: torch_cuda_${{ matrix.config.backend }}_reports
           path: reports
-          
+      - name: Generate Report and Notify Channel
+        if: always()
+        run: |
+          pip install slack_sdk tabulate
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+
   run_nightly_pipeline_level_quantization_tests:
     name: Torch quantization nightly tests
     strategy:
@@ -484,117 +574,12 @@ jobs:
         with:
           name: torch_cuda_pipeline_level_quant_reports
           path: reports
-
-  run_flax_tpu_tests:
-    name: Nightly Flax TPU Tests
-    runs-on:
-      group: gcp-ct5lp-hightpu-8t
-    if: github.event_name == 'schedule'
-
-    container:
-      image: diffusers/diffusers-flax-tpu
-      options: --shm-size "16gb" --ipc host --privileged ${{ vars.V5_LITEPOD_8_ENV}} -v /mnt/hf_cache:/mnt/hf_cache
-    defaults:
-      run:
-        shell: bash
-    steps:
-    - name: Checkout diffusers
-      uses: actions/checkout@v3
-      with:
-        fetch-depth: 2
-
-    - name: Install dependencies
-      run: |
-        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
-        python -m uv pip install -e [quality,test]
-        pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
-        python -m uv pip install pytest-reportlog
-
-    - name: Environment
-      run: python utils/print_env.py
-
-    - name: Run nightly Flax TPU tests
-      env:
-        HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }}
-      run: |
-        python -m pytest -n 0 \
-          -s -v -k "Flax" \
-          --make-reports=tests_flax_tpu \
-          --report-log=tests_flax_tpu.log \
-          tests/
-
-    - name: Failure short reports
-      if: ${{ failure() }}
-      run: |
-        cat reports/tests_flax_tpu_stats.txt
-        cat reports/tests_flax_tpu_failures_short.txt
-
-    - name: Test suite reports artifacts
-      if: ${{ always() }}
-      uses: actions/upload-artifact@v4
-      with:
-        name: flax_tpu_test_reports
-        path: reports
-
-  generate_consolidated_report:
-    name: Generate Consolidated Test Report
-    needs: [
-      run_nightly_tests_for_torch_pipelines,
-      run_nightly_tests_for_other_torch_modules,
-      run_torch_compile_tests,
-      run_big_gpu_torch_tests,
-      run_nightly_quantization_tests,
-      run_nightly_pipeline_level_quantization_tests,
-      run_nightly_onnx_tests,
-      torch_minimum_version_cuda_tests,
-      run_flax_tpu_tests
-    ]
-    if: always()
-    runs-on:
-      group: aws-general-8-plus
-    container:
-      image: diffusers/diffusers-pytorch-cpu
-    steps:
-      - name: Checkout diffusers
-        uses: actions/checkout@v3
-        with:
-          fetch-depth: 2
-
-      - name: Create reports directory
-        run: mkdir -p combined_reports
-
-      - name: Download all test reports
-        uses: actions/download-artifact@v4
-        with:
-          path: artifacts
-
-      - name: Prepare reports
+      - name: Generate Report and Notify Channel
+        if: always()
         run: |
-          # Move all report files to a single directory for processing
-          find artifacts -name "*.txt" -exec cp {} combined_reports/ \;
-
-      - name: Install dependencies
-        run: |
-          pip install -e .[test]
           pip install slack_sdk tabulate
-
-      - name: Generate consolidated report
-        run: |
-          python utils/consolidated_test_report.py \
-            --reports_dir combined_reports \
-            --output_file $CONSOLIDATED_REPORT_PATH \
-            --slack_channel_name diffusers-ci-nightly
-
-      - name: Show consolidated report
-        run: |
-          cat $CONSOLIDATED_REPORT_PATH >> $GITHUB_STEP_SUMMARY
-
-      - name: Upload consolidated report
-        uses: actions/upload-artifact@v4
-        with:
-          name: consolidated_test_report
-          path: ${{ env.CONSOLIDATED_REPORT_PATH }}
-
+          python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
+  
 # M1 runner currently not well supported
 # TODO: (Dhruv) add these back when we setup better testing for Apple Silicon
 #  run_nightly_tests_apple_m1:

.github/workflows/pr_tests.yml

@@ -291,8 +291,8 @@ jobs:
     - name: Failure short reports
       if: ${{ failure() }}
       run: |
-        cat reports/tests_peft_main_failures_short.txt
-        cat reports/tests_models_lora_peft_main_failures_short.txt
+        cat reports/tests_lora_failures_short.txt
+        cat reports/tests_models_lora_failures_short.txt
 
     - name: Test suite reports artifacts
       if: ${{ always() }}

docs/source/en/_toctree.yml

@@ -92,6 +92,8 @@
     title: API Reference
   title: Hybrid Inference
 - sections:
+  - local: using-diffusers/cogvideox
+    title: CogVideoX
   - local: using-diffusers/consisid
     title: ConsisID
   - local: using-diffusers/sdxl
@@ -176,10 +178,10 @@
 - sections:
   - local: optimization/fp16
     title: Accelerate inference
-  - local: optimization/cache
-    title: Caching
   - local: optimization/memory
     title: Reduce memory usage
+  - local: optimization/torch2.0
+    title: PyTorch 2.0
   - local: optimization/xformers
     title: xFormers
   - local: optimization/tome

docs/source/en/api/cache.md

@@ -11,19 +11,71 @@ specific language governing permissions and limitations under the License. -->
 
 # Caching methods
 
-Cache methods speedup diffusion transformers by storing and reusing intermediate outputs of specific layers, such as attention and feedforward layers, instead of recalculating them at each inference step.
+## Pyramid Attention Broadcast
 
-## CacheMixin
+[Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) from Xuanlei Zhao, Xiaolong Jin, Kai Wang, Yang You.
+
+Pyramid Attention Broadcast (PAB) is a method that speeds up inference in diffusion models by systematically skipping attention computations between successive inference steps and reusing cached attention states. The attention states are not very different between successive inference steps. The most prominent difference is in the spatial attention blocks, not as much in the temporal attention blocks, and finally the least in the cross attention blocks. Therefore, many cross attention computation blocks can be skipped, followed by the temporal and spatial attention blocks. By combining other techniques like sequence parallelism and classifier-free guidance parallelism, PAB achieves near real-time video generation.
+
+Enable PAB with [`~PyramidAttentionBroadcastConfig`] on any pipeline. For some benchmarks, refer to [this](https://github.com/huggingface/diffusers/pull/9562) pull request.
+
+```python
+import torch
+from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig
+
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+# Increasing the value of `spatial_attention_timestep_skip_range[0]` or decreasing the value of
+# `spatial_attention_timestep_skip_range[1]` will decrease the interval in which pyramid attention
+# broadcast is active, leader to slower inference speeds. However, large intervals can lead to
+# poorer quality of generated videos.
+config = PyramidAttentionBroadcastConfig(
+    spatial_attention_block_skip_range=2,
+    spatial_attention_timestep_skip_range=(100, 800),
+    current_timestep_callback=lambda: pipe.current_timestep,
+)
+pipe.transformer.enable_cache(config)
+```
+
+## Faster Cache
+
+[FasterCache](https://huggingface.co/papers/2410.19355) from Zhengyao Lv, Chenyang Si, Junhao Song, Zhenyu Yang, Yu Qiao, Ziwei Liu, Kwan-Yee K. Wong.
+
+FasterCache is a method that speeds up inference in diffusion transformers by:
+- Reusing attention states between successive inference steps, due to high similarity between them
+- Skipping unconditional branch prediction used in classifier-free guidance by revealing redundancies between unconditional and conditional branch outputs for the same timestep, and therefore approximating the unconditional branch output using the conditional branch output
+
+```python
+import torch
+from diffusers import CogVideoXPipeline, FasterCacheConfig
+
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+config = FasterCacheConfig(
+    spatial_attention_block_skip_range=2,
+    spatial_attention_timestep_skip_range=(-1, 681),
+    current_timestep_callback=lambda: pipe.current_timestep,
+    attention_weight_callback=lambda _: 0.3,
+    unconditional_batch_skip_range=5,
+    unconditional_batch_timestep_skip_range=(-1, 781),
+    tensor_format="BFCHW",
+)
+pipe.transformer.enable_cache(config)
+```
+
+### CacheMixin
 
 [[autodoc]] CacheMixin
 
-## PyramidAttentionBroadcastConfig
+### PyramidAttentionBroadcastConfig
 
 [[autodoc]] PyramidAttentionBroadcastConfig
 
 [[autodoc]] apply_pyramid_attention_broadcast
 
-## FasterCacheConfig
+### FasterCacheConfig
 
 [[autodoc]] FasterCacheConfig
 

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

@@ -98,8 +98,4 @@ To learn more about how to load LoRA weights, see the [LoRA](../../using-diffuse
 
 ## LoraBaseMixin
 
-[[autodoc]] loaders.lora_base.LoraBaseMixin
-
-## WanLoraLoaderMixin
-
-[[autodoc]] loaders.lora_pipeline.WanLoraLoaderMixin
+[[autodoc]] loaders.lora_base.LoraBaseMixin

docs/source/en/api/pipelines/cogvideox.md

@@ -13,181 +13,150 @@
 # limitations under the License.
 -->
 
-<div style="float: right;">
-  <div class="flex flex-wrap space-x-1">
-    <a href="https://huggingface.co/docs/diffusers/main/en/tutorials/using_peft_for_inference" target="_blank" rel="noopener">
-      <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
-    </a>
-  </div>
-</div>
-
 # CogVideoX
 
-[CogVideoX](https://huggingface.co/papers/2408.06072) is a large diffusion transformer model - available in 2B and 5B parameters - designed to generate longer and more consistent videos from text. This model uses a 3D causal variational autoencoder to more efficiently process video data by reducing sequence length (and associated training compute) and preventing flickering in generated videos. An "expert" transformer with adaptive LayerNorm improves alignment between text and video, and 3D full attention helps accurately capture motion and time in generated videos.
+<div class="flex flex-wrap space-x-1">
+  <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
+</div>
 
-You can find all the original CogVideoX checkpoints under the [CogVideoX](https://huggingface.co/collections/THUDM/cogvideo-66c08e62f1685a3ade464cce) collection.
+[CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer](https://huggingface.co/papers/2408.06072) from Tsinghua University & ZhipuAI, by Zhuoyi Yang, Jiayan Teng, Wendi Zheng, Ming Ding, Shiyu Huang, Jiazheng Xu, Yuanming Yang, Wenyi Hong, Xiaohan Zhang, Guanyu Feng, Da Yin, Xiaotao Gu, Yuxuan Zhang, Weihan Wang, Yean Cheng, Ting Liu, Bin Xu, Yuxiao Dong, Jie Tang.
 
-> [!TIP]
-> Click on the CogVideoX models in the right sidebar for more examples of other video generation tasks.
+The abstract from the paper is:
 
-The example below demonstrates how to generate a video optimized for memory or inference speed.
+*We introduce CogVideoX, a large-scale diffusion transformer model designed for generating videos based on text prompts. To efficently model video data, we propose to levearge a 3D Variational Autoencoder (VAE) to compresses videos along both spatial and temporal dimensions. To improve the text-video alignment, we propose an expert transformer with the expert adaptive LayerNorm to facilitate the deep fusion between the two modalities. By employing a progressive training technique, CogVideoX is adept at producing coherent, long-duration videos characterized by significant motion. In addition, we develop an effectively text-video data processing pipeline that includes various data preprocessing strategies and a video captioning method. It significantly helps enhance the performance of CogVideoX, improving both generation quality and semantic alignment. Results show that CogVideoX demonstrates state-of-the-art performance across both multiple machine metrics and human evaluations. The model weight of CogVideoX-2B is publicly available at https://github.com/THUDM/CogVideo.*
 
-<hfoptions id="usage">
-<hfoption id="memory">
+<Tip>
 
-Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving techniques.
+Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines.
 
-The quantized CogVideoX 5B model below requires ~16GB of VRAM.
+</Tip>
 
-```py
+This pipeline was contributed by [zRzRzRzRzRzRzR](https://github.com/zRzRzRzRzRzRzR). The original codebase can be found [here](https://huggingface.co/THUDM). The original weights can be found under [hf.co/THUDM](https://huggingface.co/THUDM).
+
+There are three official CogVideoX checkpoints for text-to-video and video-to-video.
+
+| checkpoints | recommended inference dtype |
+|:---:|:---:|
+| [`THUDM/CogVideoX-2b`](https://huggingface.co/THUDM/CogVideoX-2b) | torch.float16 |
+| [`THUDM/CogVideoX-5b`](https://huggingface.co/THUDM/CogVideoX-5b) | torch.bfloat16 |
+| [`THUDM/CogVideoX1.5-5b`](https://huggingface.co/THUDM/CogVideoX1.5-5b) | torch.bfloat16 |
+
+There are two official CogVideoX checkpoints available for image-to-video.
+
+| checkpoints | recommended inference dtype |
+|:---:|:---:|
+| [`THUDM/CogVideoX-5b-I2V`](https://huggingface.co/THUDM/CogVideoX-5b-I2V) | torch.bfloat16 |
+| [`THUDM/CogVideoX-1.5-5b-I2V`](https://huggingface.co/THUDM/CogVideoX-1.5-5b-I2V) | torch.bfloat16 |
+
+For the CogVideoX 1.5 series:
+- Text-to-video (T2V) works best at a resolution of 1360x768 because it was trained with that specific resolution.
+- Image-to-video (I2V) works for multiple resolutions. The width can vary from 768 to 1360, but the height must be 768. The height/width must be divisible by 16.
+- Both T2V and I2V models support generation with 81 and 161 frames and work best at this value. Exporting videos at 16 FPS is recommended.
+
+There are two official CogVideoX checkpoints that support pose controllable generation (by the [Alibaba-PAI](https://huggingface.co/alibaba-pai) team).
+
+| checkpoints | recommended inference dtype |
+|:---:|:---:|
+| [`alibaba-pai/CogVideoX-Fun-V1.1-2b-Pose`](https://huggingface.co/alibaba-pai/CogVideoX-Fun-V1.1-2b-Pose) | torch.bfloat16 |
+| [`alibaba-pai/CogVideoX-Fun-V1.1-5b-Pose`](https://huggingface.co/alibaba-pai/CogVideoX-Fun-V1.1-5b-Pose) | torch.bfloat16 |
+
+## Inference
+
+Use [`torch.compile`](https://huggingface.co/docs/diffusers/main/en/tutorials/fast_diffusion#torchcompile) to reduce the inference latency.
+
+First, load the pipeline:
+
+```python
 import torch
-from diffusers import CogVideoXPipeline, AutoModel
-from diffusers.quantizers import PipelineQuantizationConfig
-from diffusers.hooks import apply_group_offloading
-from diffusers.utils import export_to_video
-
-# quantize weights to int8 with torchao
-pipeline_quant_config = PipelineQuantizationConfig(
-  quant_backend="torchao",
-  quant_kwargs={"quant_type": "int8wo"},
-  components_to_quantize=["transformer"]
-)
-
-# fp8 layerwise weight-casting
-transformer = AutoModel.from_pretrained(
-    "THUDM/CogVideoX-5b",
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16
-)
-transformer.enable_layerwise_casting(
-    storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
-)
-
-pipeline = CogVideoXPipeline.from_pretrained(
-    "THUDM/CogVideoX-5b",
-    transformer=transformer,
-    quantization_config=pipeline_quant_config,
-    torch_dtype=torch.bfloat16
-)
-pipeline.to("cuda")
-
-# model-offloading
-pipeline.enable_model_cpu_offload()
-
-prompt = """
-A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. 
-The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. 
-Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, 
-with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
-"""
-
-video = pipeline(
-    prompt=prompt,
-    guidance_scale=6,
-    num_inference_steps=50
-).frames[0]
-export_to_video(video, "output.mp4", fps=8)
+from diffusers import CogVideoXPipeline, CogVideoXImageToVideoPipeline
+from diffusers.utils import export_to_video,load_image
+pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b").to("cuda") # or "THUDM/CogVideoX-2b" 
 ```
 
-</hfoption>
-<hfoption id="inference speed">
+If you are using the image-to-video pipeline, load it as follows:
 
-[Compilation](../../optimization/fp16#torchcompile) is slow the first time but subsequent calls to the pipeline are faster.
+```python
+pipe = CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V").to("cuda")
+```
 
-The average inference time with torch.compile on a 80GB A100 is 76.27 seconds compared to 96.89 seconds for an uncompiled model.
+Then change the memory layout of the pipelines `transformer` component to `torch.channels_last`:
+
+```python
+pipe.transformer.to(memory_format=torch.channels_last)
+```
+
+Compile the components and run inference:
+
+```python
+pipe.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
+
+# CogVideoX works well with long and well-described prompts
+prompt = "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical atmosphere of this unique musical performance."
+video = pipe(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
+```
+
+The [T2V benchmark](https://gist.github.com/a-r-r-o-w/5183d75e452a368fd17448fcc810bd3f) results on an 80GB A100 machine are:
+
+```
+Without torch.compile(): Average inference time: 96.89 seconds.
+With torch.compile(): Average inference time: 76.27 seconds.
+```
+
+### Memory optimization
+
+CogVideoX-2b requires about 19 GB of GPU memory to decode 49 frames (6 seconds of video at 8 FPS) with output resolution 720x480 (W x H), which makes it not possible to run on consumer GPUs or free-tier T4 Colab. The following memory optimizations could be used to reduce the memory footprint. For replication, you can refer to [this](https://gist.github.com/a-r-r-o-w/3959a03f15be5c9bd1fe545b09dfcc93) script.
+
+- `pipe.enable_model_cpu_offload()`:
+  - Without enabling cpu offloading, memory usage is `33 GB`
+  - With enabling cpu offloading, memory usage is `19 GB`
+- `pipe.enable_sequential_cpu_offload()`:
+  - Similar to `enable_model_cpu_offload` but can significantly reduce memory usage at the cost of slow inference
+  - When enabled, memory usage is under `4 GB`
+- `pipe.vae.enable_tiling()`:
+  - With enabling cpu offloading and tiling, memory usage is `11 GB`
+- `pipe.vae.enable_slicing()`
+
+## Quantization
+
+Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model.
+
+Refer to the [Quantization](../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`CogVideoXPipeline`] for inference with bitsandbytes.
 
 ```py
 import torch
-from diffusers import CogVideoXPipeline
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, CogVideoXTransformer3DModel, CogVideoXPipeline
 from diffusers.utils import export_to_video
+from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel
+
+quant_config = BitsAndBytesConfig(load_in_8bit=True)
+text_encoder_8bit = T5EncoderModel.from_pretrained(
+    "THUDM/CogVideoX-2b",
+    subfolder="text_encoder",
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
+)
+
+quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
+transformer_8bit = CogVideoXTransformer3DModel.from_pretrained(
+    "THUDM/CogVideoX-2b",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
+)
 
 pipeline = CogVideoXPipeline.from_pretrained(
     "THUDM/CogVideoX-2b",
-    torch_dtype=torch.float16
-).to("cuda")
-
-# torch.compile
-pipeline.transformer.to(memory_format=torch.channels_last)
-pipeline.transformer = torch.compile(
-    pipeline.transformer, mode="max-autotune", fullgraph=True
+    text_encoder=text_encoder_8bit,
+    transformer=transformer_8bit,
+    torch_dtype=torch.float16,
+    device_map="balanced",
 )
 
-prompt = """
-A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. 
-The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. 
-Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, 
-with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
-"""
-
-video = pipeline(
-    prompt=prompt,
-    guidance_scale=6,
-    num_inference_steps=50
-).frames[0]
-export_to_video(video, "output.mp4", fps=8)
+prompt = "A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting."
+video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
+export_to_video(video, "ship.mp4", fps=8)
 ```
 
-</hfoption>
-</hfoptions>
-
-## Notes
-
-- CogVideoX supports LoRAs with [`~loaders.CogVideoXLoraLoaderMixin.load_lora_weights`].
-
-  <details>
-  <summary>Show example code</summary>
-
-  ```py
-  import torch
-  from diffusers import CogVideoXPipeline
-  from diffusers.hooks import apply_group_offloading
-  from diffusers.utils import export_to_video
-
-  pipeline = CogVideoXPipeline.from_pretrained(
-      "THUDM/CogVideoX-5b",
-      torch_dtype=torch.bfloat16
-  )
-  pipeline.to("cuda")
-
-  # load LoRA weights
-  pipeline.load_lora_weights("finetrainers/CogVideoX-1.5-crush-smol-v0", adapter_name="crush-lora")
-  pipeline.set_adapters("crush-lora", 0.9)
-
-  # model-offloading
-  pipeline.enable_model_cpu_offload()
-
-  prompt = """
-  PIKA_CRUSH A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.
-  """
-  negative_prompt = "inconsistent motion, blurry motion, worse quality, degenerate outputs, deformed outputs"
-
-  video = pipeline(
-      prompt=prompt, 
-      negative_prompt=negative_prompt, 
-      num_frames=81, 
-      height=480,
-      width=768,
-      num_inference_steps=50
-  ).frames[0]
-  export_to_video(video, "output.mp4", fps=16)
-  ```
-
-  </details>
-
-- The text-to-video (T2V) checkpoints work best with a resolution of 1360x768 because that was the resolution it was pretrained on.
-
-- The image-to-video (I2V) checkpoints work with multiple resolutions. The width can vary from 768 to 1360, but the height must be 758. Both height and width must be divisible by 16.
-
-- Both T2V and I2V checkpoints work best with 81 and 161 frames. It is recommended to export the generated video at 16fps.
-
-- Refer to the table below to view memory usage when various memory-saving techniques are enabled.
-
-  | method | memory usage (enabled) | memory usage (disabled) |
-  |---|---|---|
-  | enable_model_cpu_offload | 19GB | 33GB |
-  | enable_sequential_cpu_offload | <4GB | ~33GB (very slow inference speed) |
-  | enable_tiling | 11GB (with enable_model_cpu_offload) | --- |
- 
 ## CogVideoXPipeline
 
 [[autodoc]] CogVideoXPipeline

docs/source/en/api/pipelines/deepfloyd_if.md

@@ -347,7 +347,7 @@ pipe.to("cuda")
 image = pipe(image=image, prompt="<prompt>", strength=0.3).images
 ```
 
-You can also use [`torch.compile`](../../optimization/fp16#torchcompile). Note that we have not exhaustively tested `torch.compile`
+You can also use [`torch.compile`](../../optimization/torch2.0). Note that we have not exhaustively tested `torch.compile`
 with IF and it might not give expected results.
 
 ```py

docs/source/en/api/pipelines/hunyuan_video.md

@@ -12,171 +12,78 @@
 # See the License for the specific language governing permissions and
 # limitations under the License. -->
 
-<div style="float: right;">
-  <div class="flex flex-wrap space-x-1">
-    <a href="https://huggingface.co/docs/diffusers/main/en/tutorials/using_peft_for_inference" target="_blank" rel="noopener">
-      <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
-    </a>
-  </div>
-</div>
-
 # HunyuanVideo
 
-[HunyuanVideo](https://huggingface.co/papers/2412.03603) is a 13B parameter diffusion transformer model designed to be competitive with closed-source video foundation models and enable wider community access. This model uses a "dual-stream to single-stream" architecture to separately process the video and text tokens first, before concatenating and feeding them to the transformer to fuse the multimodal information. A pretrained multimodal large language model (MLLM) is used as the encoder because it has better image-text alignment, better image detail description and reasoning, and it can be used as a zero-shot learner if system instructions are added to user prompts. Finally, HunyuanVideo uses a 3D causal variational autoencoder to more efficiently process video data at the original resolution and frame rate.
+<div class="flex flex-wrap space-x-1">
+  <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
+</div>
 
-You can find all the original HunyuanVideo checkpoints under the [Tencent](https://huggingface.co/tencent) organization.
+[HunyuanVideo](https://www.arxiv.org/abs/2412.03603) by Tencent.
 
-> [!TIP]
-> Click on the HunyuanVideo models in the right sidebar for more examples of video generation tasks.
->
-> The examples below use a checkpoint from [hunyuanvideo-community](https://huggingface.co/hunyuanvideo-community) because the weights are stored in a layout compatible with Diffusers.
+*Recent advancements in video generation have significantly impacted daily life for both individuals and industries. However, the leading video generation models remain closed-source, resulting in a notable performance gap between industry capabilities and those available to the public. In this report, we introduce HunyuanVideo, an innovative open-source video foundation model that demonstrates performance in video generation comparable to, or even surpassing, that of leading closed-source models. HunyuanVideo encompasses a comprehensive framework that integrates several key elements, including data curation, advanced architectural design, progressive model scaling and training, and an efficient infrastructure tailored for large-scale model training and inference. As a result, we successfully trained a video generative model with over 13 billion parameters, making it the largest among all open-source models. We conducted extensive experiments and implemented a series of targeted designs to ensure high visual quality, motion dynamics, text-video alignment, and advanced filming techniques. According to evaluations by professionals, HunyuanVideo outperforms previous state-of-the-art models, including Runway Gen-3, Luma 1.6, and three top-performing Chinese video generative models. By releasing the code for the foundation model and its applications, we aim to bridge the gap between closed-source and open-source communities. This initiative will empower individuals within the community to experiment with their ideas, fostering a more dynamic and vibrant video generation ecosystem. The code is publicly available at [this https URL](https://github.com/tencent/HunyuanVideo).*
 
-The example below demonstrates how to generate a video optimized for memory or inference speed.
+<Tip>
 
-<hfoptions id="usage">
-<hfoption id="memory">
+Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines.
 
-Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving techniques.
+</Tip>
 
-The quantized HunyuanVideo model below requires ~14GB of VRAM.
+Recommendations for inference:
+- Both text encoders should be in `torch.float16`.
+- Transformer should be in `torch.bfloat16`.
+- VAE should be in `torch.float16`.
+- `num_frames` should be of the form `4 * k + 1`, for example `49` or `129`.
+- For smaller resolution videos, try lower values of `shift` (between `2.0` to `5.0`) in the [Scheduler](https://huggingface.co/docs/diffusers/main/en/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler.shift). For larger resolution images, try higher values (between `7.0` and `12.0`). The default value is `7.0` for HunyuanVideo.
+- For more information about supported resolutions and other details, please refer to the original repository [here](https://github.com/Tencent/HunyuanVideo/).
+
+## Available models
+
+The following models are available for the [`HunyuanVideoPipeline`](text-to-video) pipeline:
+
+| Model name | Description |
+|:---|:---|
+| [`hunyuanvideo-community/HunyuanVideo`](https://huggingface.co/hunyuanvideo-community/HunyuanVideo) | Official HunyuanVideo (guidance-distilled). Performs best at multiple resolutions and frames. Performs best with `guidance_scale=6.0`, `true_cfg_scale=1.0` and without a negative prompt. |
+| [`https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-T2V`](https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-T2V) | Skywork's custom finetune of HunyuanVideo (de-distilled). Performs best with `97x544x960` resolution, `guidance_scale=1.0`, `true_cfg_scale=6.0` and a negative prompt. |
+
+The following models are available for the image-to-video pipeline:
+
+| Model name | Description |
+|:---|:---|
+| [`Skywork/SkyReels-V1-Hunyuan-I2V`](https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-I2V) | Skywork's custom finetune of HunyuanVideo (de-distilled). Performs best with `97x544x960` resolution. Performs best at `97x544x960` resolution, `guidance_scale=1.0`, `true_cfg_scale=6.0` and a negative prompt. |
+| [`hunyuanvideo-community/HunyuanVideo-I2V-33ch`](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-I2V) | Tecent's official HunyuanVideo 33-channel I2V model. Performs best at resolutions of 480, 720, 960, 1280. A higher `shift` value when initializing the scheduler is recommended (good values are between 7 and 20). |
+| [`hunyuanvideo-community/HunyuanVideo-I2V`](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-I2V) | Tecent's official HunyuanVideo 16-channel I2V model. Performs best at resolutions of 480, 720, 960, 1280. A higher `shift` value when initializing the scheduler is recommended (good values are between 7 and 20) |
+
+## Quantization
+
+Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model.
+
+Refer to the [Quantization](../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`HunyuanVideoPipeline`] for inference with bitsandbytes.
 
 ```py
 import torch
-from diffusers import AutoModel, HunyuanVideoPipeline
-from diffusers.quantizers import PipelineQuantizationConfig
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, HunyuanVideoTransformer3DModel, HunyuanVideoPipeline
 from diffusers.utils import export_to_video
 
-# quantize weights to int4 with bitsandbytes
-pipeline_quant_config = PipelineQuantizationConfig(
-    quant_backend="bitsandbytes_4bit",
-    quant_kwargs={
-      "load_in_4bit": True,
-      "bnb_4bit_quant_type": "nf4",
-      "bnb_4bit_compute_dtype": torch.bfloat16
-      },
-    components_to_quantize=["transformer"]
+quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
+transformer_8bit = HunyuanVideoTransformer3DModel.from_pretrained(
+    "hunyuanvideo-community/HunyuanVideo",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.bfloat16,
 )
 
 pipeline = HunyuanVideoPipeline.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
-    quantization_config=pipeline_quant_config,
-    torch_dtype=torch.bfloat16,
+    transformer=transformer_8bit,
+    torch_dtype=torch.float16,
+    device_map="balanced",
 )
 
-# model-offloading and tiling
-pipeline.enable_model_cpu_offload()
-pipeline.vae.enable_tiling()
-
-prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
+prompt = "A cat walks on the grass, realistic style."
 video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
-export_to_video(video, "output.mp4", fps=15)
+export_to_video(video, "cat.mp4", fps=15)
 ```
 
-</hfoption>
-<hfoption id="inference speed">
-
-[Compilation](../../optimization/fp16#torchcompile) is slow the first time but subsequent calls to the pipeline are faster.
-
-```py
-import torch
-from diffusers import AutoModel, HunyuanVideoPipeline
-from diffusers.quantizers import PipelineQuantizationConfig
-from diffusers.utils import export_to_video
-
-# quantize weights to int4 with bitsandbytes
-pipeline_quant_config = PipelineQuantizationConfig(
-    quant_backend="bitsandbytes_4bit",
-    quant_kwargs={
-      "load_in_4bit": True,
-      "bnb_4bit_quant_type": "nf4",
-      "bnb_4bit_compute_dtype": torch.bfloat16
-      },
-    components_to_quantize=["transformer"]
-)
-
-pipeline = HunyuanVideoPipeline.from_pretrained(
-    "hunyuanvideo-community/HunyuanVideo",
-    quantization_config=pipeline_quant_config,
-    torch_dtype=torch.bfloat16,
-)
-
-# model-offloading and tiling
-pipeline.enable_model_cpu_offload()
-pipeline.vae.enable_tiling()
-
-# torch.compile
-pipeline.transformer.to(memory_format=torch.channels_last)
-pipeline.transformer = torch.compile(
-    pipeline.transformer, mode="max-autotune", fullgraph=True
-)
-
-prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
-video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
-export_to_video(video, "output.mp4", fps=15)
-```
-
-</hfoption>
-</hfoptions>
-
-## Notes
-
-- HunyuanVideo supports LoRAs with [`~loaders.HunyuanVideoLoraLoaderMixin.load_lora_weights`].
-
-  <details>
-  <summary>Show example code</summary>
-
-  ```py
-  import torch
-  from diffusers import AutoModel, HunyuanVideoPipeline
-  from diffusers.quantizers import PipelineQuantizationConfig
-  from diffusers.utils import export_to_video
-
-  # quantize weights to int4 with bitsandbytes
-  pipeline_quant_config = PipelineQuantizationConfig(
-      quant_backend="bitsandbytes_4bit",
-      quant_kwargs={
-        "load_in_4bit": True,
-        "bnb_4bit_quant_type": "nf4",
-        "bnb_4bit_compute_dtype": torch.bfloat16
-        },
-      components_to_quantize=["transformer"]
-  )
-
-  pipeline = HunyuanVideoPipeline.from_pretrained(
-      "hunyuanvideo-community/HunyuanVideo",
-      quantization_config=pipeline_quant_config,
-      torch_dtype=torch.bfloat16,
-  )
-
-  # load LoRA weights
-  pipeline.load_lora_weights("https://huggingface.co/lucataco/hunyuan-steamboat-willie-10", adapter_name="steamboat-willie")
-  pipeline.set_adapters("steamboat-willie", 0.9)
-
-  # model-offloading and tiling
-  pipeline.enable_model_cpu_offload()
-  pipeline.vae.enable_tiling()
-
-  # use "In the style of SWR" to trigger the LoRA
-  prompt = """
-  In the style of SWR. A black and white animated scene featuring a fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys.
-  """
-  video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
-  export_to_video(video, "output.mp4", fps=15)
-  ```
-
-  </details>
-
-- Refer to the table below for recommended inference values.
-
-  | parameter | recommended value |
-  |---|---|
-  | text encoder dtype | `torch.float16` |
-  | transformer dtype | `torch.bfloat16` |
-  | vae dtype | `torch.float16` |
-  | `num_frames (k)` | 4 * `k` + 1 |
-
-- Try lower `shift` values (`2.0` to `5.0`) for lower resolution videos and higher `shift` values (`7.0` to `12.0`) for higher resolution images.
-
 ## HunyuanVideoPipeline
 
 [[autodoc]] HunyuanVideoPipeline

docs/source/en/api/pipelines/ltx_video.md

@@ -12,108 +12,322 @@
 # See the License for the specific language governing permissions and
 # limitations under the License. -->
 
-<div style="float: right;">
-  <div class="flex flex-wrap space-x-1">
-    <a href="https://huggingface.co/docs/diffusers/main/en/tutorials/using_peft_for_inference" target="_blank" rel="noopener">
-      <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
-    </a>
-    <img alt="MPS" src="https://img.shields.io/badge/MPS-000000?style=flat&logo=apple&logoColor=white%22">
-  </div>
+# LTX Video
+
+<div class="flex flex-wrap space-x-1">
+  <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
+  <img alt="MPS" src="https://img.shields.io/badge/MPS-000000?style=flat&logo=apple&logoColor=white%22">
 </div>
 
-# LTX-Video
+[LTX Video](https://huggingface.co/Lightricks/LTX-Video) is the first DiT-based video generation model capable of generating high-quality videos in real-time. It produces 24 FPS videos at a 768x512 resolution faster than they can be watched. Trained on a large-scale dataset of diverse videos, the model generates high-resolution videos with realistic and varied content. We provide a model for both text-to-video as well as image + text-to-video usecases.
 
-[LTX-Video](https://huggingface.co/Lightricks/LTX-Video) is a diffusion transformer designed for fast and real-time generation of high-resolution videos from text and images. The main feature of LTX-Video is the Video-VAE. The Video-VAE has a higher pixel to latent compression ratio (1:192) which enables more efficient video data processing and faster generation speed. To support and prevent finer details from being lost during generation, the Video-VAE decoder performs the latent to pixel conversion *and* the last denoising step.
+<Tip>
 
-You can find all the original LTX-Video checkpoints under the [Lightricks](https://huggingface.co/Lightricks) organization.
+Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers) to learn how to explore the tradeoff between scheduler speed and quality, and see the [reuse components across pipelines](../../using-diffusers/loading#reuse-a-pipeline) section to learn how to efficiently load the same components into multiple pipelines.
 
-> [!TIP]
-> Click on the LTX-Video models in the right sidebar for more examples of other video generation tasks.
+</Tip>
 
-The example below demonstrates how to generate a video optimized for memory or inference speed.
+Available models:
 
-<hfoptions id="usage">
-<hfoption id="memory">
+|  Model name   | Recommended dtype |
+|:-------------:|:-----------------:|
+| [`LTX Video 2B 0.9.0`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.safetensors) | `torch.bfloat16` |
+| [`LTX Video 2B 0.9.1`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.1.safetensors) | `torch.bfloat16` |
+| [`LTX Video 2B 0.9.5`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.5.safetensors) | `torch.bfloat16` |
+| [`LTX Video 13B 0.9.7`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-dev.safetensors) | `torch.bfloat16` |
+| [`LTX Video 13B 0.9.7 (distilled)`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-13b-0.9.7-distilled.safetensors) | `torch.bfloat16` |
+| [`LTX Video Spatial Upscaler 0.9.7`](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltxv-spatial-upscaler-0.9.7.safetensors) | `torch.bfloat16` |
 
-Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving techniques.
+Note: The recommended dtype is for the transformer component. The VAE and text encoders can be either `torch.float32`, `torch.bfloat16` or `torch.float16` but the recommended dtype is `torch.bfloat16` as used in the original repository.
 
-The LTX-Video model below requires ~10GB of VRAM.
+## Recommended settings for generation
 
-```py
+For the best results, it is recommended to follow the guidelines mentioned in the official LTX Video [repository](https://github.com/Lightricks/LTX-Video).
+
+- Some variants of LTX Video are guidance-distilled. For guidance-distilled models, `guidance_scale` must be set to `1.0`. For any other models, `guidance_scale` should be set higher (e.g., `5.0`) for good generation quality.
+- For variants with a timestep-aware VAE (LTXV 0.9.1 and above), it is recommended to set `decode_timestep` to `0.05` and `image_cond_noise_scale` to `0.025`.
+- For variants that support interpolation between multiple conditioning images and videos (LTXV 0.9.5 and above), it is recommended to use similar looking images/videos for the best results. High divergence between the conditionings may lead to abrupt transitions in the generated video.
+
+<!-- TODO(aryan): remove this warning when modular diffusers is ready -->
+
+<Tip warning={true}>
+
+The examples below show some recommended generation settings, but note that all features supported in the original [LTX Video repository](https://github.com/Lightricks/LTX-Video) are not supported in `diffusers` yet (for example, Spatio-temporal Guidance and CRF compression for image inputs). This will gradually be supported in the future. For the best possible generation quality, we recommend using the code from the original repository.
+
+</Tip>
+
+## Using LTX Video 13B 0.9.7
+
+LTX Video 0.9.7 comes with a spatial latent upscaler and a 13B parameter transformer. The inference involves generating a low resolution video first, which is very fast, followed by upscaling and refining the generated video.
+
+<!-- TODO(aryan): modify when official checkpoints are available -->
+
+```python
 import torch
-from diffusers import LTXPipeline, AutoModel
-from diffusers.hooks import apply_group_offloading
-from diffusers.utils import export_to_video
+from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
+from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
+from diffusers.utils import export_to_video, load_video
 
-# fp8 layerwise weight-casting
-transformer = AutoModel.from_pretrained(
-    "Lightricks/LTX-Video",
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16
-)
-transformer.enable_layerwise_casting(
-    storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
-)
+pipe = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-dev", torch_dtype=torch.bfloat16)
+pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipe.vae, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe_upsample.to("cuda")
+pipe.vae.enable_tiling()
 
-pipeline = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=transformer, torch_dtype=torch.bfloat16)
+def round_to_nearest_resolution_acceptable_by_vae(height, width):
+    height = height - (height % pipe.vae_temporal_compression_ratio)
+    width = width - (width % pipe.vae_temporal_compression_ratio)
+    return height, width
 
-# group-offloading
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
-apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
-apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
+video = load_video(
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4"
+)[:21]  # Use only the first 21 frames as conditioning
+condition1 = LTXVideoCondition(video=video, frame_index=0)
 
-prompt = """
-A woman with long brown hair and light skin smiles at another woman with long blonde hair.
-The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek.
-The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and 
-natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
-"""
+prompt = "The video depicts a winding mountain road covered in snow, with a single vehicle traveling along it. The road is flanked by steep, rocky cliffs and sparse vegetation. The landscape is characterized by rugged terrain and a river visible in the distance. The scene captures the solitude and beauty of a winter drive through a mountainous region."
 negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+expected_height, expected_width = 768, 1152
+downscale_factor = 2 / 3
+num_frames = 161
 
-video = pipeline(
+# Part 1. Generate video at smaller resolution
+# Text-only conditioning is also supported without the need to pass `conditions`
+downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
+downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)
+latents = pipe(
+    conditions=[condition1],
     prompt=prompt,
     negative_prompt=negative_prompt,
-    width=768,
-    height=512,
-    num_frames=161,
-    decode_timestep=0.03,
+    width=downscaled_width,
+    height=downscaled_height,
+    num_frames=num_frames,
+    num_inference_steps=30,
+    decode_timestep=0.05,
     decode_noise_scale=0.025,
-    num_inference_steps=50,
+    image_cond_noise_scale=0.0,
+    guidance_scale=5.0,
+    guidance_rescale=0.7,
+    generator=torch.Generator().manual_seed(0),
+    output_type="latent",
+).frames
+
+# Part 2. Upscale generated video using latent upsampler with fewer inference steps
+# The available latent upsampler upscales the height/width by 2x
+upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
+upscaled_latents = pipe_upsample(
+    latents=latents,
+    output_type="latent"
+).frames
+
+# Part 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended)
+video = pipe(
+    conditions=[condition1],
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    width=upscaled_width,
+    height=upscaled_height,
+    num_frames=num_frames,
+    denoise_strength=0.4,  # Effectively, 4 inference steps out of 10
+    num_inference_steps=10,
+    latents=upscaled_latents,
+    decode_timestep=0.05,
+    decode_noise_scale=0.025,
+    image_cond_noise_scale=0.0,
+    guidance_scale=5.0,
+    guidance_rescale=0.7,
+    generator=torch.Generator().manual_seed(0),
+    output_type="pil",
 ).frames[0]
+
+# Part 4. Downscale the video to the expected resolution
+video = [frame.resize((expected_width, expected_height)) for frame in video]
+
 export_to_video(video, "output.mp4", fps=24)
 ```
 
-</hfoption>
-<hfoption id="inference speed">
+## Using LTX Video 0.9.7 (distilled)
 
-[Compilation](../../optimization/fp16#torchcompile) is slow the first time but subsequent calls to the pipeline are faster.
+The same example as above can be used with the exception of the `guidance_scale` parameter. The model is both guidance and timestep distilled in order to speedup generation. It requires `guidance_scale` to be set to `1.0`. Additionally, to benefit from the timestep distillation, `num_inference_steps` can be set between `4` and `10` for good generation quality.
+
+Additionally, custom timesteps can also be used for conditioning the generation. The authors recommend using the following timesteps for best results:
+- Base model inference to prepare for upscaling: `[1000, 993, 987, 981, 975, 909, 725, 0.03]`
+- Upscaling: `[1000, 909, 725, 421, 0]`
+
+<details>
+<summary> Full example </summary>
+
+```python
+import torch
+from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
+from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
+from diffusers.utils import export_to_video, load_video
+
+pipe = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-distilled", torch_dtype=torch.bfloat16)
+pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipe.vae, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe_upsample.to("cuda")
+pipe.vae.enable_tiling()
+
+def round_to_nearest_resolution_acceptable_by_vae(height, width):
+    height = height - (height % pipe.vae_temporal_compression_ratio)
+    width = width - (width % pipe.vae_temporal_compression_ratio)
+    return height, width
+
+prompt = "artistic anatomical 3d render, utlra quality, human half full male body with transparent skin revealing structure instead of organs, muscular, intricate creative patterns, monochromatic with backlighting, lightning mesh, scientific concept art, blending biology with botany, surreal and ethereal quality, unreal engine 5, ray tracing, ultra realistic, 16K UHD, rich details. camera zooms out in a rotating fashion"
+negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+expected_height, expected_width = 768, 1152
+downscale_factor = 2 / 3
+num_frames = 161
+
+# Part 1. Generate video at smaller resolution
+downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
+downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)
+latents = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    width=downscaled_width,
+    height=downscaled_height,
+    num_frames=num_frames,
+    timesteps=[1000, 993, 987, 981, 975, 909, 725, 0.03],
+    decode_timestep=0.05,
+    decode_noise_scale=0.025,
+    image_cond_noise_scale=0.0,
+    guidance_scale=1.0,
+    guidance_rescale=0.7,
+    generator=torch.Generator().manual_seed(0),
+    output_type="latent",
+).frames
+
+# Part 2. Upscale generated video using latent upsampler with fewer inference steps
+# The available latent upsampler upscales the height/width by 2x
+upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
+upscaled_latents = pipe_upsample(
+    latents=latents,
+    adain_factor=1.0,
+    output_type="latent"
+).frames
+
+# Part 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended)
+video = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    width=upscaled_width,
+    height=upscaled_height,
+    num_frames=num_frames,
+    denoise_strength=0.999,  # Effectively, 4 inference steps out of 5
+    timesteps=[1000, 909, 725, 421, 0],
+    latents=upscaled_latents,
+    decode_timestep=0.05,
+    decode_noise_scale=0.025,
+    image_cond_noise_scale=0.0,
+    guidance_scale=1.0,
+    guidance_rescale=0.7,
+    generator=torch.Generator().manual_seed(0),
+    output_type="pil",
+).frames[0]
+
+# Part 4. Downscale the video to the expected resolution
+video = [frame.resize((expected_width, expected_height)) for frame in video]
+
+export_to_video(video, "output.mp4", fps=24)
+```
+
+</details>
+
+## Loading Single Files
+
+Loading the original LTX Video checkpoints is also possible with [`~ModelMixin.from_single_file`]. We recommend using `from_single_file` for the Lightricks series of models, as they plan to release multiple models in the future in the single file format.
+
+```python
+import torch
+from diffusers import AutoencoderKLLTXVideo, LTXImageToVideoPipeline, LTXVideoTransformer3DModel
+
+# `single_file_url` could also be https://huggingface.co/Lightricks/LTX-Video/ltx-video-2b-v0.9.1.safetensors
+single_file_url = "https://huggingface.co/Lightricks/LTX-Video/ltx-video-2b-v0.9.safetensors"
+transformer = LTXVideoTransformer3DModel.from_single_file(
+  single_file_url, torch_dtype=torch.bfloat16
+)
+vae = AutoencoderKLLTXVideo.from_single_file(single_file_url, torch_dtype=torch.bfloat16)
+pipe = LTXImageToVideoPipeline.from_pretrained(
+  "Lightricks/LTX-Video", transformer=transformer, vae=vae, torch_dtype=torch.bfloat16
+)
+
+# ... inference code ...
+```
+
+Alternatively, the pipeline can be used to load the weights with [`~FromSingleFileMixin.from_single_file`].
+
+```python
+import torch
+from diffusers import LTXImageToVideoPipeline
+from transformers import T5EncoderModel, T5Tokenizer
+
+single_file_url = "https://huggingface.co/Lightricks/LTX-Video/ltx-video-2b-v0.9.safetensors"
+text_encoder = T5EncoderModel.from_pretrained(
+  "Lightricks/LTX-Video", subfolder="text_encoder", torch_dtype=torch.bfloat16
+)
+tokenizer = T5Tokenizer.from_pretrained(
+  "Lightricks/LTX-Video", subfolder="tokenizer", torch_dtype=torch.bfloat16
+)
+pipe = LTXImageToVideoPipeline.from_single_file(
+  single_file_url, text_encoder=text_encoder, tokenizer=tokenizer, torch_dtype=torch.bfloat16
+)
+```
+
+Loading [LTX GGUF checkpoints](https://huggingface.co/city96/LTX-Video-gguf) are also supported:
 
 ```py
 import torch
+from diffusers.utils import export_to_video
+from diffusers import LTXPipeline, LTXVideoTransformer3DModel, GGUFQuantizationConfig
+
+ckpt_path = (
+    "https://huggingface.co/city96/LTX-Video-gguf/blob/main/ltx-video-2b-v0.9-Q3_K_S.gguf"
+)
+transformer = LTXVideoTransformer3DModel.from_single_file(
+    ckpt_path,
+    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
+    torch_dtype=torch.bfloat16,
+)
+pipe = LTXPipeline.from_pretrained(
+    "Lightricks/LTX-Video",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+)
+pipe.enable_model_cpu_offload()
+
+prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
+negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+video = pipe(
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    width=704,
+    height=480,
+    num_frames=161,
+    num_inference_steps=50,
+).frames[0]
+export_to_video(video, "output_gguf_ltx.mp4", fps=24)
+```
+
+Make sure to read the [documentation on GGUF](../../quantization/gguf) to learn more about our GGUF support.
+
+<!-- TODO(aryan): Update this when official weights are supported -->
+
+Loading and running inference with [LTX Video 0.9.1](https://huggingface.co/Lightricks/LTX-Video/blob/main/ltx-video-2b-v0.9.1.safetensors) weights.
+
+```python
+import torch
 from diffusers import LTXPipeline
 from diffusers.utils import export_to_video
 
-pipeline = LTXPipeline.from_pretrained(
-    "Lightricks/LTX-Video", torch_dtype=torch.bfloat16
-)
+pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-0.9.1-diffusers", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
 
-# torch.compile
-pipeline.transformer.to(memory_format=torch.channels_last)
-pipeline.transformer = torch.compile(
-    pipeline.transformer, mode="max-autotune", fullgraph=True
-)
-
-prompt = """
-A woman with long brown hair and light skin smiles at another woman with long blonde hair.
-The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek.
-The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and 
-natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
-"""
+prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
 negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
 
-video = pipeline(
+video = pipe(
     prompt=prompt,
     negative_prompt=negative_prompt,
     width=768,
@@ -126,264 +340,48 @@ video = pipeline(
 export_to_video(video, "output.mp4", fps=24)
 ```
 
-</hfoption>
-</hfoptions>
+Refer to [this section](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cogvideox#memory-optimization) to learn more about optimizing memory consumption.
 
-## Notes
+## Quantization
 
-- Refer to the following recommended settings for generation from the [LTX-Video](https://github.com/Lightricks/LTX-Video) repository.
+Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model.
 
-  - The recommended dtype for the transformer, VAE, and text encoder is `torch.bfloat16`. The VAE and text encoder can also be `torch.float32` or `torch.float16`.
-  - For guidance-distilled variants of LTX-Video, set `guidance_scale` to `1.0`. The `guidance_scale` for any other model should be set higher, like `5.0`, for good generation quality.
-  - For timestep-aware VAE variants (LTX-Video 0.9.1 and above), set `decode_timestep` to `0.05` and `image_cond_noise_scale` to `0.025`.
-  - For variants that support interpolation between multiple conditioning images and videos (LTX-Video 0.9.5 and above), use similar images and videos for the best results. Divergence from the conditioning inputs may lead to abrupt transitionts in the generated video.
+Refer to the [Quantization](../../quantization/overview) overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized [`LTXPipeline`] for inference with bitsandbytes.
 
-- LTX-Video 0.9.7 includes a spatial latent upscaler and a 13B parameter transformer. During inference, a low resolution video is quickly generated first and then upscaled and refined.
+```py
+import torch
+from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, LTXVideoTransformer3DModel, LTXPipeline
+from diffusers.utils import export_to_video
+from transformers import BitsAndBytesConfig as BitsAndBytesConfig, T5EncoderModel
 
-  <details>
-  <summary>Show example code</summary>
+quant_config = BitsAndBytesConfig(load_in_8bit=True)
+text_encoder_8bit = T5EncoderModel.from_pretrained(
+    "Lightricks/LTX-Video",
+    subfolder="text_encoder",
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
+)
 
-  ```py
-  import torch
-  from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
-  from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
-  from diffusers.utils import export_to_video, load_video
+quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
+transformer_8bit = LTXVideoTransformer3DModel.from_pretrained(
+    "Lightricks/LTX-Video",
+    subfolder="transformer",
+    quantization_config=quant_config,
+    torch_dtype=torch.float16,
+)
 
-  pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-dev", torch_dtype=torch.bfloat16)
-  pipeline_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipeline.vae, torch_dtype=torch.bfloat16)
-  pipeline.to("cuda")
-  pipe_upsample.to("cuda")
-  pipeline.vae.enable_tiling()
+pipeline = LTXPipeline.from_pretrained(
+    "Lightricks/LTX-Video",
+    text_encoder=text_encoder_8bit,
+    transformer=transformer_8bit,
+    torch_dtype=torch.float16,
+    device_map="balanced",
+)
 
-  def round_to_nearest_resolution_acceptable_by_vae(height, width):
-      height = height - (height % pipeline.vae_temporal_compression_ratio)
-      width = width - (width % pipeline.vae_temporal_compression_ratio)
-      return height, width
-
-  video = load_video(
-      "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cosmos/cosmos-video2world-input-vid.mp4"
-  )[:21]  # only use the first 21 frames as conditioning
-  condition1 = LTXVideoCondition(video=video, frame_index=0)
-
-  prompt = """
-  The video depicts a winding mountain road covered in snow, with a single vehicle 
-  traveling along it. The road is flanked by steep, rocky cliffs and sparse vegetation. 
-  The landscape is characterized by rugged terrain and a river visible in the distance. 
-  The scene captures the solitude and beauty of a winter drive through a mountainous region.
-  """
-  negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
-  expected_height, expected_width = 768, 1152
-  downscale_factor = 2 / 3
-  num_frames = 161
-
-  # 1. Generate video at smaller resolution
-  # Text-only conditioning is also supported without the need to pass `conditions`
-  downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
-  downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)
-  latents = pipeline(
-      conditions=[condition1],
-      prompt=prompt,
-      negative_prompt=negative_prompt,
-      width=downscaled_width,
-      height=downscaled_height,
-      num_frames=num_frames,
-      num_inference_steps=30,
-      decode_timestep=0.05,
-      decode_noise_scale=0.025,
-      image_cond_noise_scale=0.0,
-      guidance_scale=5.0,
-      guidance_rescale=0.7,
-      generator=torch.Generator().manual_seed(0),
-      output_type="latent",
-  ).frames
-
-  # 2. Upscale generated video using latent upsampler with fewer inference steps
-  # The available latent upsampler upscales the height/width by 2x
-  upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
-  upscaled_latents = pipe_upsample(
-      latents=latents,
-      output_type="latent"
-  ).frames
-
-  # 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended)
-  video = pipeline(
-      conditions=[condition1],
-      prompt=prompt,
-      negative_prompt=negative_prompt,
-      width=upscaled_width,
-      height=upscaled_height,
-      num_frames=num_frames,
-      denoise_strength=0.4,  # Effectively, 4 inference steps out of 10
-      num_inference_steps=10,
-      latents=upscaled_latents,
-      decode_timestep=0.05,
-      decode_noise_scale=0.025,
-      image_cond_noise_scale=0.0,
-      guidance_scale=5.0,
-      guidance_rescale=0.7,
-      generator=torch.Generator().manual_seed(0),
-      output_type="pil",
-  ).frames[0]
-
-  # 4. Downscale the video to the expected resolution
-  video = [frame.resize((expected_width, expected_height)) for frame in video]
-
-  export_to_video(video, "output.mp4", fps=24)
-  ```
-
-  </details>
-
-- LTX-Video 0.9.7 distilled model is guidance and timestep-distilled to speedup generation. It requires `guidance_scale` to be set to `1.0` and `num_inference_steps` should be set between `4` and `10` for good generation quality. You should also use the following custom timesteps for the best results.
-
-  - Base model inference to prepare for upscaling: `[1000, 993, 987, 981, 975, 909, 725, 0.03]`.
-  - Upscaling: `[1000, 909, 725, 421, 0]`.
-
-  <details>
-  <summary>Show example code</summary>
-
-  ```py
-  import torch
-  from diffusers import LTXConditionPipeline, LTXLatentUpsamplePipeline
-  from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
-  from diffusers.utils import export_to_video, load_video
-
-  pipeline = LTXConditionPipeline.from_pretrained("Lightricks/LTX-Video-0.9.7-distilled", torch_dtype=torch.bfloat16)
-  pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipeline.vae, torch_dtype=torch.bfloat16)
-  pipeline.to("cuda")
-  pipe_upsample.to("cuda")
-  pipeline.vae.enable_tiling()
-
-  def round_to_nearest_resolution_acceptable_by_vae(height, width):
-      height = height - (height % pipeline.vae_temporal_compression_ratio)
-      width = width - (width % pipeline.vae_temporal_compression_ratio)
-      return height, width
-
-  prompt = """
-  artistic anatomical 3d render, utlra quality, human half full male body with transparent 
-  skin revealing structure instead of organs, muscular, intricate creative patterns, 
-  monochromatic with backlighting, lightning mesh, scientific concept art, blending biology 
-  with botany, surreal and ethereal quality, unreal engine 5, ray tracing, ultra realistic, 
-  16K UHD, rich details. camera zooms out in a rotating fashion
-  """
-  negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
-  expected_height, expected_width = 768, 1152
-  downscale_factor = 2 / 3
-  num_frames = 161
-
-  # 1. Generate video at smaller resolution
-  downscaled_height, downscaled_width = int(expected_height * downscale_factor), int(expected_width * downscale_factor)
-  downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)
-  latents = pipeline(
-      prompt=prompt,
-      negative_prompt=negative_prompt,
-      width=downscaled_width,
-      height=downscaled_height,
-      num_frames=num_frames,
-      timesteps=[1000, 993, 987, 981, 975, 909, 725, 0.03],
-      decode_timestep=0.05,
-      decode_noise_scale=0.025,
-      image_cond_noise_scale=0.0,
-      guidance_scale=1.0,
-      guidance_rescale=0.7,
-      generator=torch.Generator().manual_seed(0),
-      output_type="latent",
-  ).frames
-
-  # 2. Upscale generated video using latent upsampler with fewer inference steps
-  # The available latent upsampler upscales the height/width by 2x
-  upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
-  upscaled_latents = pipe_upsample(
-      latents=latents,
-      adain_factor=1.0,
-      output_type="latent"
-  ).frames
-
-  # 3. Denoise the upscaled video with few steps to improve texture (optional, but recommended)
-  video = pipeline(
-      prompt=prompt,
-      negative_prompt=negative_prompt,
-      width=upscaled_width,
-      height=upscaled_height,
-      num_frames=num_frames,
-      denoise_strength=0.999,  # Effectively, 4 inference steps out of 5
-      timesteps=[1000, 909, 725, 421, 0],
-      latents=upscaled_latents,
-      decode_timestep=0.05,
-      decode_noise_scale=0.025,
-      image_cond_noise_scale=0.0,
-      guidance_scale=1.0,
-      guidance_rescale=0.7,
-      generator=torch.Generator().manual_seed(0),
-      output_type="pil",
-  ).frames[0]
-
-  # 4. Downscale the video to the expected resolution
-  video = [frame.resize((expected_width, expected_height)) for frame in video]
-
-  export_to_video(video, "output.mp4", fps=24)
-  ```
-
-  </details>
-
-- LTX-Video supports LoRAs with [`~loaders.LTXVideoLoraLoaderMixin.load_lora_weights`].
-
-  <details>
-  <summary>Show example code</summary>
-
-  ```py
-  import torch
-  from diffusers import LTXConditionPipeline
-  from diffusers.utils import export_to_video, load_image
-
-  pipeline = LTXConditionPipeline.from_pretrained(
-      "Lightricks/LTX-Video-0.9.5", torch_dtype=torch.bfloat16
-  )
-
-  pipeline.load_lora_weights("Lightricks/LTX-Video-Cakeify-LoRA", adapter_name="cakeify")
-  pipeline.set_adapters("cakeify")
-
-  # use "CAKEIFY" to trigger the LoRA
-  prompt = "CAKEIFY a person using a knife to cut a cake shaped like a Pikachu plushie"
-  image = load_image("https://huggingface.co/Lightricks/LTX-Video-Cakeify-LoRA/resolve/main/assets/images/pikachu.png")
-
-  video = pipeline(
-      prompt=prompt,
-      image=image,
-      width=576,
-      height=576,
-      num_frames=161,
-      decode_timestep=0.03,
-      decode_noise_scale=0.025,
-      num_inference_steps=50,
-  ).frames[0]
-  export_to_video(video, "output.mp4", fps=26)
-  ```
-
-  </details>
-
-- LTX-Video supports loading from single files, such as [GGUF checkpoints](../../quantization/gguf), with [`loaders.FromOriginalModelMixin.from_single_file`] or [`loaders.FromSingleFileMixin.from_single_file`].
-
-  <details>
-  <summary>Show example code</summary>
-
-  ```py
-  import torch
-  from diffusers.utils import export_to_video
-  from diffusers import LTXPipeline, AutoModel, GGUFQuantizationConfig
-
-  transformer = AutoModel.from_single_file(
-      "https://huggingface.co/city96/LTX-Video-gguf/blob/main/ltx-video-2b-v0.9-Q3_K_S.gguf"
-      quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
-      torch_dtype=torch.bfloat16
-  )
-  pipeline = LTXPipeline.from_pretrained(
-      "Lightricks/LTX-Video",
-      transformer=transformer,
-      torch_dtype=torch.bfloat16
-  )
-  ```
-
-  </details>
+prompt = "A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting."
+video = pipeline(prompt=prompt, num_frames=161, num_inference_steps=50).frames[0]
+export_to_video(video, "ship.mp4", fps=24)
+```
 
 ## LTXPipeline
 

docs/source/en/api/pipelines/sana_sprint.md

@@ -88,46 +88,12 @@ image.save("sana.png")
 
 Users can tweak the `max_timesteps` value for experimenting with the visual quality of the generated outputs. The default `max_timesteps` value was obtained with an inference-time search process. For more details about it, check out the paper.
 
-## Image to Image 
-
-The [`SanaSprintImg2ImgPipeline`] is a pipeline for image-to-image generation. It takes an input image and a prompt, and generates a new image based on the input image and the prompt.
-
-```py
-import torch
-from diffusers import SanaSprintImg2ImgPipeline
-from diffusers.utils.loading_utils import load_image
-
-image = load_image(
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
-)
-
-pipe = SanaSprintImg2ImgPipeline.from_pretrained(
-    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", 
-    torch_dtype=torch.bfloat16)
-pipe.to("cuda")
-
-image = pipe(
-    prompt="a cute pink bear", 
-    image=image, 
-    strength=0.5, 
-    height=832, 
-    width=480
-).images[0]
-image.save("output.png")
-```
-
 ## SanaSprintPipeline
 
 [[autodoc]] SanaSprintPipeline
   - all
   - __call__
 
-## SanaSprintImg2ImgPipeline
-
-[[autodoc]] SanaSprintImg2ImgPipeline
-  - all
-  - __call__
-
 
 ## SanaPipelineOutput
 

docs/source/en/api/pipelines/wan.md

@@ -12,170 +12,128 @@
 # See the License for the specific language governing permissions and
 # limitations under the License. -->
 
-<div style="float: right;">
-  <div class="flex flex-wrap space-x-1">
-    <a href="https://huggingface.co/docs/diffusers/main/en/tutorials/using_peft_for_inference" target="_blank" rel="noopener">
-      <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
-    </a>
-  </div>
+# Wan
+
+<div class="flex flex-wrap space-x-1">
+  <img alt="LoRA" src="https://img.shields.io/badge/LoRA-d8b4fe?style=flat"/>
 </div>
 
-# Wan2.1
+[Wan 2.1](https://github.com/Wan-Video/Wan2.1) by the Alibaba Wan Team.
 
-[Wan-2.1](https://huggingface.co/papers/2503.20314) by the Wan Team.
+<!-- TODO(aryan): update abstract once paper is out -->
 
-*This report presents Wan, a comprehensive and open suite of video foundation models designed to push the boundaries of video generation. Built upon the mainstream diffusion transformer paradigm, Wan achieves significant advancements in generative capabilities through a series of innovations, including our novel VAE, scalable pre-training strategies, large-scale data curation, and automated evaluation metrics. These contributions collectively enhance the model's performance and versatility. Specifically, Wan is characterized by four key features: Leading Performance: The 14B model of Wan, trained on a vast dataset comprising billions of images and videos, demonstrates the scaling laws of video generation with respect to both data and model size. It consistently outperforms the existing open-source models as well as state-of-the-art commercial solutions across multiple internal and external benchmarks, demonstrating a clear and significant performance superiority. Comprehensiveness: Wan offers two capable models, i.e., 1.3B and 14B parameters, for efficiency and effectiveness respectively. It also covers multiple downstream applications, including image-to-video, instruction-guided video editing, and personal video generation, encompassing up to eight tasks. Consumer-Grade Efficiency: The 1.3B model demonstrates exceptional resource efficiency, requiring only 8.19 GB VRAM, making it compatible with a wide range of consumer-grade GPUs. Openness: We open-source the entire series of Wan, including source code and all models, with the goal of fostering the growth of the video generation community. This openness seeks to significantly expand the creative possibilities of video production in the industry and provide academia with high-quality video foundation models. All the code and models are available at [this https URL](https://github.com/Wan-Video/Wan2.1).*
+## Generating Videos with Wan 2.1
 
-You can find all the original Wan2.1 checkpoints under the [Wan-AI](https://huggingface.co/Wan-AI) organization.
+We will first need to install some additional dependencies.
 
-The following Wan models are supported in Diffusers:
-- [Wan 2.1 T2V 1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers)
-- [Wan 2.1 T2V 14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B-Diffusers)
-- [Wan 2.1 I2V 14B - 480P](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers)
-- [Wan 2.1 I2V 14B - 720P](https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P-Diffusers)
-- [Wan 2.1 FLF2V 14B - 720P](https://huggingface.co/Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers)
-- [Wan 2.1 VACE 1.3B](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers)
-- [Wan 2.1 VACE 14B](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers)
-
-> [!TIP]
-> Click on the Wan2.1 models in the right sidebar for more examples of video generation.
-
-### Text-to-Video Generation
-
-The example below demonstrates how to generate a video from text optimized for memory or inference speed.
-
-<hfoptions id="T2V usage">
-<hfoption id="T2V memory">
-
-Refer to the [Reduce memory usage](../../optimization/memory) guide for more details about the various memory saving techniques.
-
-The Wan2.1 text-to-video model below requires ~13GB of VRAM.
-
-```py
-# pip install ftfy
-import torch
-import numpy as np
-from diffusers import AutoModel, WanPipeline
-from diffusers.quantizers import PipelineQuantizationConfig
-from diffusers.hooks.group_offloading import apply_group_offloading
-from diffusers.utils import export_to_video, load_image
-from transformers import UMT5EncoderModel
-
-text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
-vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
-transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
-
-# group-offloading
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-apply_group_offloading(text_encoder,
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="block_level",
-    num_blocks_per_group=4
-)
-transformer.enable_group_offload(
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="leaf_level",
-    use_stream=True
-)
-
-pipeline = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-    vae=vae,
-    transformer=transformer,
-    text_encoder=text_encoder,
-    torch_dtype=torch.bfloat16
-)
-pipeline.to("cuda")
-
-prompt = """
-The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
-shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
-negative_prompt = """
-Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, 
-low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, 
-misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
-"""
-
-output = pipeline(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    num_frames=81,
-    guidance_scale=5.0,
-).frames[0]
-export_to_video(output, "output.mp4", fps=16)
+```shell
+pip install -u ftfy imageio-ffmpeg imageio
 ```
 
-</hfoption>
-<hfoption id="T2V inference speed">
+### Text to Video Generation
 
-[Compilation](../../optimization/fp16#torchcompile) is slow the first time but subsequent calls to the pipeline are faster.
+The following example requires 11GB VRAM to run and uses the smaller `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` model. You can switch it out
+for the larger `Wan2.1-I2V-14B-720P-Diffusers` or `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` if you have at least 35GB VRAM available.
 
-```py
-# pip install ftfy
-import torch
-import numpy as np
-from diffusers import AutoModel, WanPipeline
-from diffusers.hooks.group_offloading import apply_group_offloading
-from diffusers.utils import export_to_video, load_image
-from transformers import UMT5EncoderModel
+```python
+from diffusers import WanPipeline
+from diffusers.utils import export_to_video
 
-text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
-vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
-transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
+# Available models: Wan-AI/Wan2.1-I2V-14B-720P-Diffusers or Wan-AI/Wan2.1-I2V-14B-480P-Diffusers
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 
-pipeline = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-    vae=vae,
-    transformer=transformer,
-    text_encoder=text_encoder,
-    torch_dtype=torch.bfloat16
-)
-pipeline.to("cuda")
+pipe = WanPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()
 
-# torch.compile
-pipeline.transformer.to(memory_format=torch.channels_last)
-pipeline.transformer = torch.compile(
-    pipeline.transformer, mode="max-autotune", fullgraph=True
-)
+prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+num_frames = 33
 
-prompt = """
-The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
-shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
-negative_prompt = """
-Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, 
-low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, 
-misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
-"""
-
-output = pipeline(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    num_frames=81,
-    guidance_scale=5.0,
-).frames[0]
-export_to_video(output, "output.mp4", fps=16)
+frames = pipe(prompt=prompt, negative_prompt=negative_prompt, num_frames=num_frames).frames[0]
+export_to_video(frames, "wan-t2v.mp4", fps=16)
 ```
 
-</hfoption>
-</hfoptions>
+<Tip>
+You can improve the quality of the generated video by running the decoding step in full precision.
+</Tip>
 
-### First-Last-Frame-to-Video Generation
+```python
+from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.utils import export_to_video
 
-The example below demonstrates how to use the image-to-video pipeline to generate a video using a text description, a starting frame, and an ending frame.
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 
-<hfoptions id="FLF2V usage">
-<hfoption id="usage">
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+
+# replace this with pipe.to("cuda") if you have sufficient VRAM
+pipe.enable_model_cpu_offload()
+
+prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+num_frames = 33
+
+frames = pipe(prompt=prompt, num_frames=num_frames).frames[0]
+export_to_video(frames, "wan-t2v.mp4", fps=16)
+```
+
+### Image to Video Generation
+
+The Image to Video pipeline requires loading the `AutoencoderKLWan` and the `CLIPVisionModel` components in full precision. The following example will need at least
+35GB of VRAM to run.
+
+```python
+import torch
+import numpy as np
+from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
+from diffusers.utils import export_to_video, load_image
+from transformers import CLIPVisionModel
+
+# Available models: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers, Wan-AI/Wan2.1-I2V-14B-720P-Diffusers
+model_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=torch.float32
+)
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
+)
+
+# replace this with pipe.to("cuda") if you have sufficient VRAM
+pipe.enable_model_cpu_offload()
+
+image = load_image(
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+)
+
+max_area = 480 * 832
+aspect_ratio = image.height / image.width
+mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+image = image.resize((width, height))
+
+prompt = (
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in "
+    "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+)
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+
+num_frames = 33
+
+output = pipe(
+    image=image,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=height,
+    width=width,
+    num_frames=num_frames,
+    guidance_scale=5.0,
+).frames[0]
+export_to_video(output, "wan-i2v.mp4", fps=16)
+```
+
+### First and Last Frame Interpolation
 
 ```python
 import numpy as np
@@ -208,13 +166,13 @@ def aspect_ratio_resize(image, pipe, max_area=720 * 1280):
 def center_crop_resize(image, height, width):
     # Calculate resize ratio to match first frame dimensions
     resize_ratio = max(width / image.width, height / image.height)
-
+    
     # Resize the image
     width = round(image.width * resize_ratio)
     height = round(image.height * resize_ratio)
     size = [width, height]
     image = TF.center_crop(image, size)
-
+    
     return image, height, width
 
 first_frame, height, width = aspect_ratio_resize(first_frame, pipe)
@@ -229,103 +187,320 @@ output = pipe(
 export_to_video(output, "output.mp4", fps=16)
 ```
 
-</hfoption>
-</hfoptions>
+### Video to Video Generation
 
-### Any-to-Video Controllable Generation
+```python
+import torch
+from diffusers.utils import load_video, export_to_video
+from diffusers import AutoencoderKLWan, WanVideoToVideoPipeline, UniPCMultistepScheduler
 
-Wan VACE supports various generation techniques which achieve controllable video generation. Some of the capabilities include:
-- Control to Video (Depth, Pose, Sketch, Flow, Grayscale, Scribble, Layout, Boundary Box, etc.). Recommended library for preprocessing videos to obtain control videos: [huggingface/controlnet_aux]()
-- Image/Video to Video (first frame, last frame, starting clip, ending clip, random clips)
-- Inpainting and Outpainting
-- Subject to Video (faces, object, characters, etc.)
-- Composition to Video (reference anything, animate anything, swap anything, expand anything, move anything, etc.)
+# Available models: Wan-AI/Wan2.1-T2V-14B-Diffusers, Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+vae = AutoencoderKLWan.from_pretrained(
+    model_id, subfolder="vae", torch_dtype=torch.float32
+)
+pipe = WanVideoToVideoPipeline.from_pretrained(
+    model_id, vae=vae, torch_dtype=torch.bfloat16
+)
+flow_shift = 3.0  # 5.0 for 720P, 3.0 for 480P
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=flow_shift
+)
+# change to pipe.to("cuda") if you have sufficient VRAM
+pipe.enable_model_cpu_offload()
 
-The code snippets available in [this](https://github.com/huggingface/diffusers/pull/11582) pull request demonstrate some examples of how videos can be generated with controllability signals.
+prompt = "A robot standing on a mountain top. The sun is setting in the background"
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+video = load_video(
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hiker.mp4"
+)
+output = pipe(
+    video=video,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=480,
+    width=512,
+    guidance_scale=7.0,
+    strength=0.7,
+).frames[0]
 
-The general rule of thumb to keep in mind when preparing inputs for the VACE pipeline is that the input images, or frames of a video that you want to use for conditioning, should have a corresponding mask that is black in color. The black mask signifies that the model will not generate new content for that area, and only use those parts for conditioning the generation process. For parts/frames that should be generated by the model, the mask should be white in color.
+export_to_video(output, "wan-v2v.mp4", fps=16)
+```
 
-## Notes
+## Memory Optimizations for Wan 2.1
 
-- Wan2.1 supports LoRAs with [`~loaders.WanLoraLoaderMixin.load_lora_weights`].
+Base inference with the large 14B Wan 2.1 models can take up to 35GB of VRAM when generating videos at 720p resolution. We'll outline a few memory optimizations we can apply to reduce the VRAM required to run the model.
 
-  <details>
-  <summary>Show example code</summary>
+We'll use `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` model in these examples to demonstrate the memory savings, but the techniques are applicable to all model checkpoints.
 
-  ```py
-  # pip install ftfy
-  import torch
-  from diffusers import AutoModel, WanPipeline
-  from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-  from diffusers.utils import export_to_video
+### Group Offloading the Transformer and UMT5 Text Encoder
 
-  vae = AutoModel.from_pretrained(
-      "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32
-  )
-  pipeline = WanPipeline.from_pretrained(
-      "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", vae=vae, torch_dtype=torch.bfloat16
-  )
-  pipeline.scheduler = UniPCMultistepScheduler.from_config(
-      pipeline.scheduler.config, flow_shift=5.0
-  )
-  pipeline.to("cuda")
+Find more information about group offloading [here](../optimization/memory.md)
 
-  pipeline.load_lora_weights("benjamin-paine/steamboat-willie-1.3b", adapter_name="steamboat-willie")
-  pipeline.set_adapters("steamboat-willie")
+#### Block Level Group Offloading
 
-  pipeline.enable_model_cpu_offload()
+We can reduce our VRAM requirements by applying group offloading to the larger model components of the pipeline; the `WanTransformer3DModel` and `UMT5EncoderModel`. Group offloading will break up the individual modules of a model and offload/onload them onto your GPU as needed during inference. In this example, we'll apply `block_level` offloading, which will group the modules in a model into blocks of size `num_blocks_per_group` and offload/onload them to GPU. Moving to between CPU and GPU does add latency to the inference process. You can trade off between latency and memory savings by increasing or decreasing the `num_blocks_per_group`.
 
-  # use "steamboat willie style" to trigger the LoRA
-  prompt = """
-  steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot, 
-  revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-  for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-  Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
-  shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-  """
+The following example will now only require 14GB of VRAM to run, but will take approximately 30 minutes to generate a video.
 
-  output = pipeline(
-      prompt=prompt,
-      num_frames=81,
-      guidance_scale=5.0,
-  ).frames[0]
-  export_to_video(output, "output.mp4", fps=16)
-  ```
+```python
+import torch
+import numpy as np
+from diffusers import AutoencoderKLWan, WanTransformer3DModel, WanImageToVideoPipeline
+from diffusers.hooks.group_offloading import apply_group_offloading
+from diffusers.utils import export_to_video, load_image
+from transformers import UMT5EncoderModel, CLIPVisionModel
 
-  </details>
+# Available models: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers, Wan-AI/Wan2.1-I2V-14B-720P-Diffusers
+model_id = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=torch.float32
+)
 
-- [`WanTransformer3DModel`] and [`AutoencoderKLWan`] supports loading from single files with [`~loaders.FromSingleFileMixin.from_single_file`].
+text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
 
-  <details>
-  <summary>Show example code</summary>
+onload_device = torch.device("cuda")
+offload_device = torch.device("cpu")
 
-  ```py
-  # pip install ftfy
-  import torch
-  from diffusers import WanPipeline, AutoModel
+apply_group_offloading(text_encoder,
+    onload_device=onload_device,
+    offload_device=offload_device,
+    offload_type="block_level",
+    num_blocks_per_group=4
+)
 
-  vae = AutoModel.from_single_file(
-      "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors"
-  )
-  transformer = AutoModel.from_single_file(
-      "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors",
-      torch_dtype=torch.bfloat16
-  )
-  pipeline = WanPipeline.from_pretrained(
-      "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-      vae=vae,
-      transformer=transformer,
-      torch_dtype=torch.bfloat16
-  )
-  ```
+transformer.enable_group_offload(
+    onload_device=onload_device,
+    offload_device=offload_device,
+    offload_type="block_level",
+    num_blocks_per_group=4,
+)
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id,
+    vae=vae,
+    transformer=transformer,
+    text_encoder=text_encoder,
+    image_encoder=image_encoder,
+    torch_dtype=torch.bfloat16
+)
+# Since we've offloaded the larger models already, we can move the rest of the model components to GPU
+pipe.to("cuda")
 
-  </details>
+image = load_image(
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+)
 
-- Set the [`AutoencoderKLWan`] dtype to `torch.float32` for better decoding quality.
+max_area = 720 * 832
+aspect_ratio = image.height / image.width
+mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+image = image.resize((width, height))
 
-- The number of frames per second (fps) or `k` should be calculated by `4 * k + 1`.
+prompt = (
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in "
+    "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+)
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
 
-- Try lower `shift` values (`2.0` to `5.0`) for lower resolution videos and higher `shift` values (`7.0` to `12.0`) for higher resolution images.
+num_frames = 33
+
+output = pipe(
+    image=image,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=height,
+    width=width,
+    num_frames=num_frames,
+    guidance_scale=5.0,
+).frames[0]
+
+export_to_video(output, "wan-i2v.mp4", fps=16)
+```
+
+#### Block Level Group Offloading with CUDA Streams
+
+We can speed up group offloading inference, by enabling the use of [CUDA streams](https://pytorch.org/docs/stable/generated/torch.cuda.Stream.html). However, using CUDA streams requires moving the model parameters into pinned memory. This allocation is handled by Pytorch under the hood, and can result in a significant spike in CPU RAM usage. Please consider this option if your CPU RAM is atleast 2X the size of the model you are group offloading.
+
+In the following example we will use CUDA streams when group offloading the `WanTransformer3DModel`. When testing on an A100, this example will require 14GB of VRAM, 52GB of CPU RAM, but will generate a video in approximately 9 minutes.
+
+```python
+import torch
+import numpy as np
+from diffusers import AutoencoderKLWan, WanTransformer3DModel, WanImageToVideoPipeline
+from diffusers.hooks.group_offloading import apply_group_offloading
+from diffusers.utils import export_to_video, load_image
+from transformers import UMT5EncoderModel, CLIPVisionModel
+
+# Available models: Wan-AI/Wan2.1-I2V-14B-480P-Diffusers, Wan-AI/Wan2.1-I2V-14B-720P-Diffusers
+model_id = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=torch.float32
+)
+
+text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
+
+onload_device = torch.device("cuda")
+offload_device = torch.device("cpu")
+
+apply_group_offloading(text_encoder,
+    onload_device=onload_device,
+    offload_device=offload_device,
+    offload_type="block_level",
+    num_blocks_per_group=4
+)
+
+transformer.enable_group_offload(
+    onload_device=onload_device,
+    offload_device=offload_device,
+    offload_type="leaf_level",
+    use_stream=True
+)
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id,
+    vae=vae,
+    transformer=transformer,
+    text_encoder=text_encoder,
+    image_encoder=image_encoder,
+    torch_dtype=torch.bfloat16
+)
+# Since we've offloaded the larger models already, we can move the rest of the model components to GPU
+pipe.to("cuda")
+
+image = load_image(
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+)
+
+max_area = 720 * 832
+aspect_ratio = image.height / image.width
+mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+image = image.resize((width, height))
+
+prompt = (
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in "
+    "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+)
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+
+num_frames = 33
+
+output = pipe(
+    image=image,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=height,
+    width=width,
+    num_frames=num_frames,
+    guidance_scale=5.0,
+).frames[0]
+
+export_to_video(output, "wan-i2v.mp4", fps=16)
+```
+
+### Applying Layerwise Casting to the Transformer
+
+Find more information about layerwise casting [here](../optimization/memory.md)
+
+In this example, we will model offloading with layerwise casting. Layerwise casting will downcast each layer's weights to `torch.float8_e4m3fn`, temporarily upcast to `torch.bfloat16` during the forward pass of the layer, then revert to `torch.float8_e4m3fn` afterward. This approach reduces memory requirements by approximately 50% while introducing a minor quality reduction in the generated video due to the precision trade-off.
+
+This example will require 20GB of VRAM.
+
+```python
+import torch
+import numpy as np
+from diffusers import AutoencoderKLWan, WanTransformer3DModel, WanImageToVideoPipeline
+from diffusers.hooks.group_offloading import apply_group_offloading
+from diffusers.utils import export_to_video, load_image
+from transformers import UMT5EncoderModel, CLIPVisionModel
+
+model_id = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
+image_encoder = CLIPVisionModel.from_pretrained(
+    model_id, subfolder="image_encoder", torch_dtype=torch.float32
+)
+text_encoder = UMT5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16)
+vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+
+transformer = WanTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16)
+transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)
+
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id,
+    vae=vae,
+    transformer=transformer,
+    text_encoder=text_encoder,
+    image_encoder=image_encoder,
+    torch_dtype=torch.bfloat16
+)
+pipe.enable_model_cpu_offload()
+image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg")
+
+max_area = 720 * 832
+aspect_ratio = image.height / image.width
+mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]
+height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+image = image.resize((width, height))
+prompt = (
+    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in "
+    "the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+)
+negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+num_frames = 33
+
+output = pipe(
+    image=image,
+    prompt=prompt,
+    negative_prompt=negative_prompt,
+    height=height,
+    width=width,
+    num_frames=num_frames,
+    num_inference_steps=50,
+    guidance_scale=5.0,
+).frames[0]
+export_to_video(output, "wan-i2v.mp4", fps=16)
+```
+
+## Using a Custom Scheduler
+
+Wan can be used with many different schedulers, each with their own benefits regarding speed and generation quality. By default, Wan uses the `UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)` scheduler. You can use a different scheduler as follows:
+
+```python
+from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler, WanPipeline
+
+scheduler_a = FlowMatchEulerDiscreteScheduler(shift=5.0)
+scheduler_b = UniPCMultistepScheduler(prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=4.0)
+
+pipe = WanPipeline.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", scheduler=<CUSTOM_SCHEDULER_HERE>)
+
+# or,
+pipe.scheduler = <CUSTOM_SCHEDULER_HERE>
+```
+
+## Using Single File Loading with Wan 2.1
+
+The `WanTransformer3DModel` and `AutoencoderKLWan` models support loading checkpoints in their original format via the `from_single_file` loading
+method.
+
+```python
+import torch
+from diffusers import WanPipeline, WanTransformer3DModel
+
+ckpt_path = "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors"
+transformer = WanTransformer3DModel.from_single_file(ckpt_path, torch_dtype=torch.bfloat16)
+
+pipe = WanPipeline.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers", transformer=transformer)
+```
+
+## Recommendations for Inference
+- Keep `AutencoderKLWan` in `torch.float32` for better decoding quality.
+- `num_frames` should satisfy the following constraint: `(num_frames - 1) % 4 == 0`
+- For smaller resolution videos, try lower values of `shift` (between `2.0` to `5.0`) in the [Scheduler](https://huggingface.co/docs/diffusers/main/en/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler.shift). For larger resolution videos, try higher values (between `7.0` and `12.0`). The default value is `3.0` for Wan.
 
 ## WanPipeline
 
@@ -339,18 +514,6 @@ The general rule of thumb to keep in mind when preparing inputs for the VACE pip
   - all
   - __call__
 
-## WanVACEPipeline
-
-[[autodoc]] WanVACEPipeline
-  - all
-  - __call__
-
-## WanVideoToVideoPipeline
-
-[[autodoc]] WanVideoToVideoPipeline
-  - all
-  - __call__
-
 ## WanPipelineOutput
 
-[[autodoc]] pipelines.wan.pipeline_output.WanPipelineOutput
+[[autodoc]] pipelines.wan.pipeline_output.WanPipelineOutput

docs/source/en/optimization/cache.md

@@ -1,69 +0,0 @@
-<!-- Copyright 2024 The HuggingFace Team. All rights reserved.
-
-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. -->
-
-# Caching
-
-Caching accelerates inference by storing and reusing intermediate outputs of different layers, such as attention and feedforward layers, instead of performing the entire computation at each inference step. It significantly improves generation speed at the expense of more memory and doesn't require additional training.
-
-This guide shows you how to use the caching methods supported in Diffusers.
-
-## Pyramid Attention Broadcast
-
-[Pyramid Attention Broadcast (PAB)](https://huggingface.co/papers/2408.12588) is based on the observation that attention outputs aren't that different between successive timesteps of the generation process. The attention differences are smallest in the cross attention layers and are generally cached over a longer timestep range. This is followed by temporal attention and spatial attention layers.
-
-> [!TIP]
-> Not all video models have three types of attention (cross, temporal, and spatial)!
-
-PAB can be combined with other techniques like sequence parallelism and classifier-free guidance parallelism (data parallelism) for near real-time video generation.
-
-Set up and pass a [`PyramidAttentionBroadcastConfig`] to a pipeline's transformer to enable it. The `spatial_attention_block_skip_range` controls how often to skip attention calculations in the spatial attention blocks and the `spatial_attention_timestep_skip_range` is the range of timesteps to skip. Take care to choose an appropriate range because a smaller interval can lead to slower inference speeds and a larger interval can result in lower generation quality.
-
-```python
-import torch
-from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig
-
-pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
-pipeline.to("cuda")
-
-config = PyramidAttentionBroadcastConfig(
-    spatial_attention_block_skip_range=2,
-    spatial_attention_timestep_skip_range=(100, 800),
-    current_timestep_callback=lambda: pipe.current_timestep,
-)
-pipeline.transformer.enable_cache(config)
-```
-
-## FasterCache
-
-[FasterCache](https://huggingface.co/papers/2410.19355) caches and reuses attention features similar to [PAB](#pyramid-attention-broadcast) since output differences are small for each successive timestep.
-
-This method may also choose to skip the unconditional branch prediction, when using classifier-free guidance for sampling (common in most base models), and estimate it from the conditional branch prediction if there is significant redundancy in the predicted latent outputs between successive timesteps.
-
-Set up and pass a [`FasterCacheConfig`] to a pipeline's transformer to enable it.
-
-```python
-import torch
-from diffusers import CogVideoXPipeline, FasterCacheConfig
-
-pipe line= CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
-pipeline.to("cuda")
-
-config = FasterCacheConfig(
-    spatial_attention_block_skip_range=2,
-    spatial_attention_timestep_skip_range=(-1, 681),
-    current_timestep_callback=lambda: pipe.current_timestep,
-    attention_weight_callback=lambda _: 0.3,
-    unconditional_batch_skip_range=5,
-    unconditional_batch_timestep_skip_range=(-1, 781),
-    tensor_format="BFCHW",
-)
-pipeline.transformer.enable_cache(config)
-```

docs/source/en/optimization/fp16.md

@@ -150,24 +150,6 @@ pipeline(prompt, num_inference_steps=30).images[0]
 
 Compilation is slow the first time, but once compiled, it is significantly faster. Try to only use the compiled pipeline on the same type of inference operations. Calling the compiled pipeline on a different image size retriggers compilation which is slow and inefficient.
 
-### Regional compilation
-
-[Regional compilation](https://docs.pytorch.org/tutorials/recipes/regional_compilation.html) reduces the cold start compilation time by only compiling a specific repeated region (or block) of the model instead of the entire model. The compiler reuses the cached and compiled code for the other blocks.
-
-[Accelerate](https://huggingface.co/docs/accelerate/index) provides the [compile_regions](https://github.com/huggingface/accelerate/blob/273799c85d849a1954a4f2e65767216eb37fa089/src/accelerate/utils/other.py#L78) method for automatically compiling the repeated blocks of a `nn.Module` sequentially. The rest of the model is compiled separately.
-
-```py
-# pip install -U accelerate
-import torch
-from diffusers import StableDiffusionXLPipeline
-from accelerate.utils import compile regions
-
-pipeline = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
-).to("cuda")
-pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)
-```
-
 ### Graph breaks
 
 It is important to specify `fullgraph=True` in torch.compile to ensure there are no graph breaks in the underlying model. This allows you to take advantage of torch.compile without any performance degradation. For the UNet and VAE, this changes how you access the return variables.
@@ -188,12 +170,6 @@ The `step()` function is [called](https://github.com/huggingface/diffusers/blob/
 
 In general, the `sigmas` should [stay on the CPU](https://github.com/huggingface/diffusers/blob/35a969d297cba69110d175ee79c59312b9f49e1e/src/diffusers/schedulers/scheduling_euler_discrete.py#L240) to avoid the communication sync and latency.
 
-### Benchmarks
-
-Refer to the [diffusers/benchmarks](https://huggingface.co/datasets/diffusers/benchmarks) dataset to see inference latency and memory usage data for compiled pipelines.
-
-The [diffusers-torchao](https://github.com/sayakpaul/diffusers-torchao#benchmarking-results) repository also contains benchmarking results for compiled versions of Flux and CogVideoX.
-
 ## Dynamic quantization
 
 [Dynamic quantization](https://pytorch.org/tutorials/recipes/recipes/dynamic_quantization.html) improves inference speed by reducing precision to enable faster math operations. This particular type of quantization determines how to scale the activations based on the data at runtime rather than using a fixed scaling factor. As a result, the scaling factor is more accurately aligned with the data.

docs/source/en/optimization/tome.md

@@ -93,4 +93,4 @@ To reproduce this benchmark, feel free to use this [script](https://gist.github.
 |          |                |              2 |         OOM |             13 |               10.78 |
 |          |                |              1 |         OOM |           6.66 |                5.54 |
 
-As seen in the tables above, the speed-up from `tomesd` becomes more pronounced for larger image resolutions. It is also interesting to note that with `tomesd`, it is possible to run the pipeline on a higher resolution like 1024x1024. You may be able to speed-up inference even more with [`torch.compile`](fp16#torchcompile).
+As seen in the tables above, the speed-up from `tomesd` becomes more pronounced for larger image resolutions. It is also interesting to note that with `tomesd`, it is possible to run the pipeline on a higher resolution like 1024x1024. You may be able to speed-up inference even more with [`torch.compile`](torch2.0).

docs/source/en/optimization/torch2.0.md

@@ -0,0 +1,438 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+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.
+-->
+
+# PyTorch 2.0
+
+🤗 Diffusers supports the latest optimizations from [PyTorch 2.0](https://pytorch.org/get-started/pytorch-2.0/) which include:
+
+1. A memory-efficient attention implementation, scaled dot product attention, without requiring any extra dependencies such as xFormers.
+2. [`torch.compile`](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html), a just-in-time (JIT) compiler to provide an extra performance boost when individual models are compiled.
+
+Both of these optimizations require PyTorch 2.0 or later and 🤗 Diffusers > 0.13.0.
+
+```bash
+pip install --upgrade torch diffusers
+```
+
+## Scaled dot product attention
+
+[`torch.nn.functional.scaled_dot_product_attention`](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention) (SDPA) is an optimized and memory-efficient attention (similar to xFormers) that automatically enables several other optimizations depending on the model inputs and GPU type. SDPA is enabled by default if you're using PyTorch 2.0 and the latest version of 🤗 Diffusers, so you don't need to add anything to your code.
+
+However, if you want to explicitly enable it, you can set a [`DiffusionPipeline`] to use [`~models.attention_processor.AttnProcessor2_0`]:
+
+```diff
+  import torch
+  from diffusers import DiffusionPipeline
++ from diffusers.models.attention_processor import AttnProcessor2_0
+
+  pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
++ pipe.unet.set_attn_processor(AttnProcessor2_0())
+
+  prompt = "a photo of an astronaut riding a horse on mars"
+  image = pipe(prompt).images[0]
+```
+
+SDPA should be as fast and memory efficient as `xFormers`; check the [benchmark](#benchmark) for more details.
+
+In some cases - such as making the pipeline more deterministic or converting it to other formats - it may be helpful to use the vanilla attention processor, [`~models.attention_processor.AttnProcessor`]. To revert to [`~models.attention_processor.AttnProcessor`], call the [`~UNet2DConditionModel.set_default_attn_processor`] function on the pipeline:
+
+```diff
+  import torch
+  from diffusers import DiffusionPipeline
+
+  pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
++ pipe.unet.set_default_attn_processor()
+
+  prompt = "a photo of an astronaut riding a horse on mars"
+  image = pipe(prompt).images[0]
+```
+
+## torch.compile
+
+The `torch.compile` function can often provide an additional speed-up to your PyTorch code. In 🤗 Diffusers, it is usually best to wrap the UNet with `torch.compile` because it does most of the heavy lifting in the pipeline.
+
+```python
+from diffusers import DiffusionPipeline
+import torch
+
+pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True).to("cuda")
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+images = pipe(prompt, num_inference_steps=steps, num_images_per_prompt=batch_size).images[0]
+```
+
+Depending on GPU type, `torch.compile` can provide an *additional speed-up* of **5-300x** on top of SDPA! If you're using more recent GPU architectures such as Ampere (A100, 3090), Ada (4090), and Hopper (H100), `torch.compile` is able to squeeze even more performance out of these GPUs.
+
+Compilation requires some time to complete, so it is best suited for situations where you prepare your pipeline once and then perform the same type of inference operations multiple times. For example, calling the compiled pipeline on a different image size triggers compilation again which can be expensive.
+
+For more information and different options about `torch.compile`, refer to the [`torch_compile`](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) tutorial.
+
+> [!TIP]
+> Learn more about other ways PyTorch 2.0 can help optimize your model in the [Accelerate inference of text-to-image diffusion models](../tutorials/fast_diffusion) tutorial.
+
+### Regional compilation
+
+Compiling the whole model usually has a big problem space for optimization. Models are often composed of multiple repeated blocks. [Regional compilation](https://pytorch.org/tutorials/recipes/regional_compilation.html) compiles the repeated block first (a transformer encoder block, for example), so that the Torch compiler would re-use its cached/optimized generated code for the other blocks, reducing (often massively) the cold start compilation time observed on the first inference call.
+
+Enabling regional compilation might require simple yet intrusive changes to the
+modeling code. However, 🤗 Accelerate provides a utility [`compile_regions()`](https://huggingface.co/docs/accelerate/main/en/usage_guides/compilation#how-to-use-regional-compilation) which automatically compiles
+the repeated blocks of the provided `nn.Module` sequentially, and the rest of the model separately. This helps with reducing cold start time while keeping most (if not all) of the speedup you would get from full compilation.
+
+```py
+# Make sure you're on the latest `accelerate`: `pip install -U accelerate`.
+from accelerate.utils import compile_regions
+
+pipe.unet = compile_regions(pipe.unet, mode="reduce-overhead", fullgraph=True)
+```
+
+As you may have noticed `compile_regions()` takes the same arguments as `torch.compile()`, allowing flexibility.
+
+## Benchmark
+
+We conducted a comprehensive benchmark with PyTorch 2.0's efficient attention implementation and `torch.compile` across different GPUs and batch sizes for five of our most used pipelines. The code is benchmarked on 🤗 Diffusers v0.17.0.dev0 to optimize `torch.compile` usage (see [here](https://github.com/huggingface/diffusers/pull/3313) for more details).
+
+Expand the dropdown below to find the code used to benchmark each pipeline:
+
+<details>
+
+### Stable Diffusion text-to-image
+
+```python
+from diffusers import DiffusionPipeline
+import torch
+
+path = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+
+run_compile = True  # Set True / False
+
+pipe = DiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, use_safetensors=True)
+pipe = pipe.to("cuda")
+pipe.unet.to(memory_format=torch.channels_last)
+
+if run_compile:
+    print("Run torch compile")
+    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+prompt = "ghibli style, a fantasy landscape with castles"
+
+for _ in range(3):
+    images = pipe(prompt=prompt).images
+```
+
+### Stable Diffusion image-to-image
+
+```python
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers.utils import load_image
+import torch
+
+url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+init_image = load_image(url)
+init_image = init_image.resize((512, 512))
+
+path = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+
+run_compile = True  # Set True / False
+
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(path, torch_dtype=torch.float16, use_safetensors=True)
+pipe = pipe.to("cuda")
+pipe.unet.to(memory_format=torch.channels_last)
+
+if run_compile:
+    print("Run torch compile")
+    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+prompt = "ghibli style, a fantasy landscape with castles"
+
+for _ in range(3):
+    image = pipe(prompt=prompt, image=init_image).images[0]
+```
+
+### Stable Diffusion inpainting
+
+```python
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers.utils import load_image
+import torch
+
+img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+
+init_image = load_image(img_url).resize((512, 512))
+mask_image = load_image(mask_url).resize((512, 512))
+
+path = "runwayml/stable-diffusion-inpainting"
+
+run_compile = True  # Set True / False
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(path, torch_dtype=torch.float16, use_safetensors=True)
+pipe = pipe.to("cuda")
+pipe.unet.to(memory_format=torch.channels_last)
+
+if run_compile:
+    print("Run torch compile")
+    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+
+prompt = "ghibli style, a fantasy landscape with castles"
+
+for _ in range(3):
+    image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
+```
+
+### ControlNet
+
+```python
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers.utils import load_image
+import torch
+
+url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+init_image = load_image(url)
+init_image = init_image.resize((512, 512))
+
+path = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+
+run_compile = True  # Set True / False
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16, use_safetensors=True)
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    path, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
+)
+
+pipe = pipe.to("cuda")
+pipe.unet.to(memory_format=torch.channels_last)
+pipe.controlnet.to(memory_format=torch.channels_last)
+
+if run_compile:
+    print("Run torch compile")
+    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+    pipe.controlnet = torch.compile(pipe.controlnet, mode="reduce-overhead", fullgraph=True)
+
+prompt = "ghibli style, a fantasy landscape with castles"
+
+for _ in range(3):
+    image = pipe(prompt=prompt, image=init_image).images[0]
+```
+
+### DeepFloyd IF text-to-image + upscaling
+
+```python
+from diffusers import DiffusionPipeline
+import torch
+
+run_compile = True  # Set True / False
+
+pipe_1 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16, use_safetensors=True)
+pipe_1.to("cuda")
+pipe_2 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-II-M-v1.0", variant="fp16", text_encoder=None, torch_dtype=torch.float16, use_safetensors=True)
+pipe_2.to("cuda")
+pipe_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16, use_safetensors=True)
+pipe_3.to("cuda")
+
+
+pipe_1.unet.to(memory_format=torch.channels_last)
+pipe_2.unet.to(memory_format=torch.channels_last)
+pipe_3.unet.to(memory_format=torch.channels_last)
+
+if run_compile:
+    pipe_1.unet = torch.compile(pipe_1.unet, mode="reduce-overhead", fullgraph=True)
+    pipe_2.unet = torch.compile(pipe_2.unet, mode="reduce-overhead", fullgraph=True)
+    pipe_3.unet = torch.compile(pipe_3.unet, mode="reduce-overhead", fullgraph=True)
+
+prompt = "the blue hulk"
+
+prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16)
+neg_prompt_embeds = torch.randn((1, 2, 4096), dtype=torch.float16)
+
+for _ in range(3):
+    image_1 = pipe_1(prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images
+    image_2 = pipe_2(image=image_1, prompt_embeds=prompt_embeds, negative_prompt_embeds=neg_prompt_embeds, output_type="pt").images
+    image_3 = pipe_3(prompt=prompt, image=image_1, noise_level=100).images
+```
+</details>
+
+The graph below highlights the relative speed-ups for the [`StableDiffusionPipeline`] across five GPU families with PyTorch 2.0 and `torch.compile` enabled. The benchmarks for the following graphs are measured in *number of iterations/second*.
+
+![t2i_speedup](https://huggingface.co/datasets/diffusers/docs-images/resolve/main/pt2_benchmarks/t2i_speedup.png)
+
+To give you an even better idea of how this speed-up holds for the other pipelines, consider the following
+graph for an A100 with PyTorch 2.0 and `torch.compile`:
+
+![a100_numbers](https://huggingface.co/datasets/diffusers/docs-images/resolve/main/pt2_benchmarks/a100_numbers.png)
+
+In the following tables, we report our findings in terms of the *number of iterations/second*.
+
+### A100 (batch size: 1)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 21.66 | 23.13 | 44.03 | 49.74 |
+| SD - img2img | 21.81 | 22.40 | 43.92 | 46.32 |
+| SD - inpaint | 22.24 | 23.23 | 43.76 | 49.25 |
+| SD - controlnet | 15.02 | 15.82 | 32.13 | 36.08 |
+| IF | 20.21 / <br>13.84 / <br>24.00 | 20.12 / <br>13.70 / <br>24.03 | ❌ | 97.34 / <br>27.23 / <br>111.66 |
+| SDXL - txt2img | 8.64 | 9.9 | - | - |
+
+### A100 (batch size: 4)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 11.6 | 13.12 | 14.62 | 17.27 |
+| SD - img2img | 11.47 | 13.06 | 14.66 | 17.25 |
+| SD - inpaint | 11.67 | 13.31 | 14.88 | 17.48 |
+| SD - controlnet | 8.28 | 9.38 | 10.51 | 12.41 |
+| IF | 25.02 | 18.04 | ❌ | 48.47 |
+| SDXL - txt2img | 2.44 | 2.74 | - | - |
+
+### A100 (batch size: 16)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 3.04 | 3.6 | 3.83 | 4.68 |
+| SD - img2img | 2.98 | 3.58 | 3.83 | 4.67 |
+| SD - inpaint | 3.04 | 3.66 | 3.9 | 4.76 |
+| SD - controlnet | 2.15 | 2.58 | 2.74 | 3.35 |
+| IF | 8.78 | 9.82 | ❌ | 16.77 |
+| SDXL - txt2img | 0.64 | 0.72 | - | - |
+
+### V100 (batch size: 1)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 18.99 | 19.14 | 20.95 | 22.17 |
+| SD - img2img | 18.56 | 19.18 | 20.95 | 22.11 |
+| SD - inpaint | 19.14 | 19.06 | 21.08 | 22.20 |
+| SD - controlnet | 13.48 | 13.93 | 15.18 | 15.88 |
+| IF |  20.01 / <br>9.08 / <br>23.34 | 19.79 / <br>8.98 / <br>24.10 | ❌ | 55.75 / <br>11.57 / <br>57.67 |
+
+### V100 (batch size: 4)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 5.96 | 5.89 | 6.83 | 6.86 |
+| SD - img2img | 5.90 | 5.91 | 6.81 | 6.82 |
+| SD - inpaint | 5.99 | 6.03 | 6.93 | 6.95 |
+| SD - controlnet | 4.26 | 4.29 | 4.92 | 4.93 |
+| IF | 15.41 | 14.76 | ❌ | 22.95 |
+
+### V100 (batch size: 16)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 1.66 | 1.66 | 1.92 | 1.90 |
+| SD - img2img | 1.65 | 1.65 | 1.91 | 1.89 |
+| SD - inpaint | 1.69 | 1.69 | 1.95 | 1.93 |
+| SD - controlnet | 1.19 | 1.19 | OOM after warmup | 1.36 |
+| IF | 5.43 | 5.29 | ❌ | 7.06 |
+
+### T4 (batch size: 1)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 6.9 | 6.95 | 7.3 | 7.56 |
+| SD - img2img | 6.84 | 6.99 | 7.04 | 7.55 |
+| SD - inpaint | 6.91 | 6.7 | 7.01 | 7.37 |
+| SD - controlnet | 4.89 | 4.86 | 5.35 | 5.48 |
+| IF | 17.42 / <br>2.47 / <br>18.52 | 16.96 / <br>2.45 / <br>18.69 | ❌ | 24.63 / <br>2.47 / <br>23.39 |
+| SDXL - txt2img | 1.15 | 1.16 | - | - |
+
+### T4 (batch size: 4)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 1.79 | 1.79 | 2.03 | 1.99 |
+| SD - img2img | 1.77 | 1.77 | 2.05 | 2.04 |
+| SD - inpaint | 1.81 | 1.82 | 2.09 | 2.09 |
+| SD - controlnet | 1.34 | 1.27 | 1.47 | 1.46 |
+| IF | 5.79 |  5.61 | ❌ | 7.39 |
+| SDXL - txt2img | 0.288 | 0.289 | - | - |
+
+### T4 (batch size: 16)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 2.34s | 2.30s | OOM after 2nd iteration | 1.99s |
+| SD - img2img | 2.35s | 2.31s | OOM after warmup | 2.00s |
+| SD - inpaint | 2.30s | 2.26s | OOM after 2nd iteration | 1.95s |
+| SD - controlnet | OOM after 2nd iteration | OOM after 2nd iteration | OOM after warmup | OOM after warmup |
+| IF * | 1.44 | 1.44 | ❌ | 1.94 |
+| SDXL - txt2img | OOM | OOM | - | - |
+
+### RTX 3090 (batch size: 1)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 22.56 | 22.84 | 23.84 | 25.69 |
+| SD - img2img | 22.25 | 22.61 | 24.1 | 25.83 |
+| SD - inpaint | 22.22 | 22.54 | 24.26 | 26.02 |
+| SD - controlnet | 16.03 | 16.33 | 17.38 | 18.56 |
+| IF | 27.08 / <br>9.07 / <br>31.23 | 26.75 / <br>8.92 / <br>31.47 | ❌ | 68.08 / <br>11.16 / <br>65.29 |
+
+### RTX 3090 (batch size: 4)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 6.46 | 6.35 | 7.29 | 7.3 |
+| SD - img2img | 6.33 | 6.27 | 7.31 | 7.26 |
+| SD - inpaint | 6.47 | 6.4 | 7.44 | 7.39 |
+| SD - controlnet | 4.59 | 4.54 | 5.27 | 5.26 |
+| IF | 16.81 | 16.62 | ❌ | 21.57 |
+
+### RTX 3090 (batch size: 16)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 1.7 | 1.69 | 1.93 | 1.91 |
+| SD - img2img | 1.68 | 1.67 | 1.93 | 1.9 |
+| SD - inpaint | 1.72 | 1.71 | 1.97 | 1.94 |
+| SD - controlnet | 1.23 | 1.22 | 1.4 | 1.38 |
+| IF | 5.01 | 5.00 | ❌ | 6.33 |
+
+### RTX 4090 (batch size: 1)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 40.5 | 41.89 | 44.65 | 49.81 |
+| SD - img2img | 40.39 | 41.95 | 44.46 | 49.8 |
+| SD - inpaint | 40.51 | 41.88 | 44.58 | 49.72 |
+| SD - controlnet | 29.27 | 30.29 | 32.26 | 36.03 |
+| IF | 69.71 / <br>18.78 / <br>85.49 | 69.13 / <br>18.80 / <br>85.56 | ❌ | 124.60 / <br>26.37 / <br>138.79 |
+| SDXL - txt2img | 6.8 | 8.18 | - | - |
+
+### RTX 4090 (batch size: 4)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 12.62 | 12.84 | 15.32 | 15.59 |
+| SD - img2img | 12.61 | 12,.79 | 15.35 | 15.66 |
+| SD - inpaint | 12.65 | 12.81 | 15.3 | 15.58 |
+| SD - controlnet | 9.1 | 9.25 | 11.03 | 11.22 |
+| IF | 31.88 | 31.14 | ❌ | 43.92 |
+| SDXL - txt2img | 2.19 | 2.35 | - | - |
+
+### RTX 4090 (batch size: 16)
+
+| **Pipeline** | **torch 2.0 - <br>no compile** | **torch nightly - <br>no compile** | **torch 2.0 - <br>compile** | **torch nightly - <br>compile** |
+|:---:|:---:|:---:|:---:|:---:|
+| SD - txt2img | 3.17 | 3.2 | 3.84 | 3.85 |
+| SD - img2img | 3.16 | 3.2 | 3.84 | 3.85 |
+| SD - inpaint | 3.17 | 3.2 | 3.85 | 3.85 |
+| SD - controlnet | 2.23 | 2.3 | 2.7 | 2.75 |
+| IF | 9.26 | 9.2 | ❌ | 13.31 |
+| SDXL - txt2img | 0.52 | 0.53 | - | - |
+
+## Notes
+
+* Follow this [PR](https://github.com/huggingface/diffusers/pull/3313) for more details on the environment used for conducting the benchmarks.
+* For the DeepFloyd IF pipeline where batch sizes > 1, we only used a batch size of > 1 in the first IF pipeline for text-to-image generation and NOT for upscaling. That means the two upscaling pipelines received a batch size of 1.
+
+*Thanks to [Horace He](https://github.com/Chillee) from the PyTorch team for their support in improving our support of `torch.compile()` in Diffusers.*

docs/source/en/quantization/torchao.md

@@ -56,7 +56,7 @@ image = pipe(
 image.save("output.png")
 ```
 
-TorchAO is fully compatible with [torch.compile](../optimization/fp16#torchcompile), setting it apart from other quantization methods. This makes it easy to speed up inference with just one line of code.
+TorchAO is fully compatible with [torch.compile](./optimization/torch2.0#torchcompile), setting it apart from other quantization methods. This makes it easy to speed up inference with just one line of code.
 
 ```python
 # In the above code, add the following after initializing the transformer
@@ -91,7 +91,7 @@ The quantization methods supported are as follows:
 
 Some quantization methods are aliases (for example, `int8wo` is the commonly used shorthand for `int8_weight_only`). This allows using the quantization methods described in the torchao docs as-is, while also making it convenient to remember their shorthand notations.
 
-Refer to the [official torchao documentation](https://docs.pytorch.org/ao/stable/index.html) for a better understanding of the available quantization methods and the exhaustive list of configuration options available.
+Refer to the official torchao documentation for a better understanding of the available quantization methods and the exhaustive list of configuration options available.
 
 ## Serializing and Deserializing quantized models
 
@@ -155,5 +155,5 @@ transformer.load_state_dict(state_dict, strict=True, assign=True)
 
 ## Resources
 
-- [TorchAO Quantization API](https://docs.pytorch.org/ao/stable/index.html)
+- [TorchAO Quantization API](https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md)
 - [Diffusers-TorchAO examples](https://github.com/sayakpaul/diffusers-torchao)

docs/source/en/stable_diffusion.md

@@ -256,6 +256,6 @@ make_image_grid(images, 2, 2)
 
 In this tutorial, you learned how to optimize a [`DiffusionPipeline`] for computational and memory efficiency as well as improving the quality of generated outputs. If you're interested in making your pipeline even faster, take a look at the following resources:
 
-- Learn how [PyTorch 2.0](./optimization/fp16) and [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) can yield 5 - 300% faster inference speed. On an A100 GPU, inference can be up to 50% faster!
+- Learn how [PyTorch 2.0](./optimization/torch2.0) and [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) can yield 5 - 300% faster inference speed. On an A100 GPU, inference can be up to 50% faster!
 - If you can't use PyTorch 2, we recommend you install [xFormers](./optimization/xformers). Its memory-efficient attention mechanism works great with PyTorch 1.13.1 for faster speed and reduced memory consumption.
 - Other optimization techniques, such as model offloading, are covered in [this guide](./optimization/fp16).

docs/source/en/training/overview.md

@@ -59,5 +59,5 @@ pip install -r requirements_sdxl.txt
 
 To speedup training and reduce memory-usage, we recommend:
 
-- using PyTorch 2.0 or higher to automatically use [scaled dot product attention](../optimization/fp16#scaled-dot-product-attention) during training (you don't need to make any changes to the training code)
+- using PyTorch 2.0 or higher to automatically use [scaled dot product attention](../optimization/torch2.0#scaled-dot-product-attention) during training (you don't need to make any changes to the training code)
 - installing [xFormers](../optimization/xformers) to enable memory-efficient attention

docs/source/en/tutorials/using_peft_for_inference.md

@@ -103,7 +103,7 @@ pipeline("A cute cnmt eating a slice of pizza, stunning color scheme, masterpiec
 
 ## torch.compile
 
-[torch.compile](../optimization/fp16#torchcompile) speeds up inference by compiling the PyTorch model to use optimized kernels. Before compiling, the LoRA weights need to be fused into the base model and unloaded first.
+[torch.compile](../optimization/torch2.0#torchcompile) speeds up inference by compiling the PyTorch model to use optimized kernels. Before compiling, the LoRA weights need to be fused into the base model and unloaded first.
 
 ```py
 import torch

docs/source/en/using-diffusers/cogvideox.md

@@ -0,0 +1,120 @@
+<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+
+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.
+-->
+# CogVideoX
+
+CogVideoX is a text-to-video generation model focused on creating more coherent videos aligned with a prompt. It achieves this using several methods.
+
+- a 3D variational autoencoder that compresses videos spatially and temporally, improving compression rate and video accuracy.
+
+- an expert transformer block to help align text and video, and a 3D full attention module for capturing and creating spatially and temporally accurate videos.
+ 
+
+
+## Load model checkpoints
+Model weights may be stored in separate subfolders on the Hub or locally, in which case, you should use the [`~DiffusionPipeline.from_pretrained`] method.
+
+
+```py
+from diffusers import CogVideoXPipeline, CogVideoXImageToVideoPipeline
+pipe = CogVideoXPipeline.from_pretrained(
+    "THUDM/CogVideoX-2b",
+    torch_dtype=torch.float16
+)
+
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V",
+    torch_dtype=torch.bfloat16
+)
+
+```
+
+## Text-to-Video
+For text-to-video, pass a text prompt. By default, CogVideoX generates a 720x480 video for the best results.
+
+```py
+import torch
+from diffusers import CogVideoXPipeline
+from diffusers.utils import export_to_video
+
+prompt = "An elderly gentleman, with a serene expression, sits at the water's edge, a steaming cup of tea by his side. He is engrossed in his artwork, brush in hand, as he renders an oil painting on a canvas that's propped up against a small, weathered table. The sea breeze whispers through his silver hair, gently billowing his loose-fitting white shirt, while the salty air adds an intangible element to his masterpiece in progress. The scene is one of tranquility and inspiration, with the artist's canvas capturing the vibrant hues of the setting sun reflecting off the tranquil sea."
+
+pipe = CogVideoXPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b",
+    torch_dtype=torch.bfloat16
+)
+
+pipe.enable_model_cpu_offload()
+pipe.vae.enable_tiling()
+
+video = pipe(
+    prompt=prompt,
+    num_videos_per_prompt=1,
+    num_inference_steps=50,
+    num_frames=49,
+    guidance_scale=6,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+).frames[0]
+
+export_to_video(video, "output.mp4", fps=8)
+
+```
+
+
+<div class="flex justify-center">
+    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_out.gif" alt="generated image of an astronaut in a jungle"/>
+</div>
+
+
+## Image-to-Video
+
+
+You'll use the [THUDM/CogVideoX-5b-I2V](https://huggingface.co/THUDM/CogVideoX-5b-I2V)  checkpoint for this guide.
+
+```py
+import torch
+from diffusers import CogVideoXImageToVideoPipeline
+from diffusers.utils import export_to_video, load_image
+
+prompt = "A vast, shimmering ocean flows gracefully under a twilight sky, its waves undulating in a mesmerizing dance of blues and greens. The surface glints with the last rays of the setting sun, casting golden highlights that ripple across the water. Seagulls soar above, their cries blending with the gentle roar of the waves. The horizon stretches infinitely, where the ocean meets the sky in a seamless blend of hues. Close-ups reveal the intricate patterns of the waves, capturing the fluidity and dynamic beauty of the sea in motion."
+image = load_image(image="cogvideox_rocket.png")
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V",
+    torch_dtype=torch.bfloat16
+)
+ 
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+
+video = pipe(
+    prompt=prompt,
+    image=image,
+    num_videos_per_prompt=1,
+    num_inference_steps=50,
+    num_frames=49,
+    guidance_scale=6,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+).frames[0]
+
+export_to_video(video, "output.mp4", fps=8)
+```
+
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_rocket.png"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">initial image</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_outrocket.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">generated video</figcaption>
+  </div>
+</div>
+

docs/source/en/using-diffusers/conditional_image_generation.md

@@ -303,7 +303,7 @@ There are many types of conditioning inputs you can use, and 🤗 Diffusers supp
 
 Diffusion models are large, and the iterative nature of denoising an image is computationally expensive and intensive. But this doesn't mean you need access to powerful - or even many - GPUs to use them. There are many optimization techniques for running diffusion models on consumer and free-tier resources. For example, you can load model weights in half-precision to save GPU memory and increase speed or offload the entire model to the GPU to save even more memory.
 
-PyTorch 2.0 also supports a more memory-efficient attention mechanism called [*scaled dot product attention*](../optimization/fp16#scaled-dot-product-attention) that is automatically enabled if you're using PyTorch 2.0. You can combine this with [`torch.compile`](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) to speed your code up even more:
+PyTorch 2.0 also supports a more memory-efficient attention mechanism called [*scaled dot product attention*](../optimization/torch2.0#scaled-dot-product-attention) that is automatically enabled if you're using PyTorch 2.0. You can combine this with [`torch.compile`](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) to speed your code up even more:
 
 ```py
 from diffusers import AutoPipelineForText2Image
@@ -313,4 +313,4 @@ pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stab
 pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
 ```
 
-For more tips on how to optimize your code to save memory and speed up inference, read the [Accelerate inference](../optimization/fp16) and [Reduce memory usage](../optimization/memory) guides.
+For more tips on how to optimize your code to save memory and speed up inference, read the [Memory and speed](../optimization/fp16) and [Torch 2.0](../optimization/torch2.0) guides.

docs/source/en/using-diffusers/img2img.md

@@ -35,7 +35,7 @@ pipeline.enable_xformers_memory_efficient_attention()
 
 <Tip>
 
-You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention).
+You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, then you don't need to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention).
 
 </Tip>
 
@@ -589,17 +589,17 @@ make_image_grid([init_image, depth_image, image_control_net, image_elden_ring],
 
 ## Optimize
 
-Running diffusion models is computationally expensive and intensive, but with a few optimization tricks, it is entirely possible to run them on consumer and free-tier GPUs. For example, you can use a more memory-efficient form of attention such as PyTorch 2.0's [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention) or [xFormers](../optimization/xformers) (you can use one or the other, but there's no need to use both). You can also offload the model to the GPU while the other pipeline components wait on the CPU.
+Running diffusion models is computationally expensive and intensive, but with a few optimization tricks, it is entirely possible to run them on consumer and free-tier GPUs. For example, you can use a more memory-efficient form of attention such as PyTorch 2.0's [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention) or [xFormers](../optimization/xformers) (you can use one or the other, but there's no need to use both). You can also offload the model to the GPU while the other pipeline components wait on the CPU.
 
 ```diff
 + pipeline.enable_model_cpu_offload()
 + pipeline.enable_xformers_memory_efficient_attention()
 ```
 
-With [`torch.compile`](../optimization/fp16#torchcompile), you can boost your inference speed even more by wrapping your UNet with it:
+With [`torch.compile`](../optimization/torch2.0#torchcompile), you can boost your inference speed even more by wrapping your UNet with it:
 
 ```py
 pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
 ```
 
-To learn more, take a look at the [Reduce memory usage](../optimization/memory) and [Accelerate inference](../optimization/fp16) guides.
+To learn more, take a look at the [Reduce memory usage](../optimization/memory) and [Torch 2.0](../optimization/torch2.0) guides.

docs/source/en/using-diffusers/inpaint.md

@@ -35,7 +35,7 @@ pipeline.enable_xformers_memory_efficient_attention()
 
 <Tip>
 
-You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, it's not necessary to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention).
+You'll notice throughout the guide, we use [`~DiffusionPipeline.enable_model_cpu_offload`] and [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`], to save memory and increase inference speed. If you're using PyTorch 2.0, it's not necessary to call [`~DiffusionPipeline.enable_xformers_memory_efficient_attention`] on your pipeline because it'll already be using PyTorch 2.0's native [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention).
 
 </Tip>
 
@@ -788,7 +788,7 @@ make_image_grid([init_image, mask_image, image, image_elden_ring], rows=2, cols=
 
 ## Optimize
 
-It can be difficult and slow to run diffusion models if you're resource constrained, but it doesn't have to be with a few optimization tricks. One of the biggest (and easiest) optimizations you can enable is switching to memory-efficient attention. If you're using PyTorch 2.0, [scaled-dot product attention](../optimization/fp16#scaled-dot-product-attention) is automatically enabled and you don't need to do anything else. For non-PyTorch 2.0 users, you can install and use [xFormers](../optimization/xformers)'s implementation of memory-efficient attention. Both options reduce memory usage and accelerate inference.
+It can be difficult and slow to run diffusion models if you're resource constrained, but it doesn't have to be with a few optimization tricks. One of the biggest (and easiest) optimizations you can enable is switching to memory-efficient attention. If you're using PyTorch 2.0, [scaled-dot product attention](../optimization/torch2.0#scaled-dot-product-attention) is automatically enabled and you don't need to do anything else. For non-PyTorch 2.0 users, you can install and use [xFormers](../optimization/xformers)'s implementation of memory-efficient attention. Both options reduce memory usage and accelerate inference.
 
 You can also offload the model to the CPU to save even more memory:
 
@@ -797,10 +797,10 @@ You can also offload the model to the CPU to save even more memory:
 + pipeline.enable_model_cpu_offload()
 ```
 
-To speed-up your inference code even more, use [`torch_compile`](../optimization/fp16#torchcompile). You should wrap `torch.compile` around the most intensive component in the pipeline which is typically the UNet:
+To speed-up your inference code even more, use [`torch_compile`](../optimization/torch2.0#torchcompile). You should wrap `torch.compile` around the most intensive component in the pipeline which is typically the UNet:
 
 ```py
 pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
 ```
 
-Learn more in the [Reduce memory usage](../optimization/memory) and [Accelerate inference](../optimization/fp16) guides.
+Learn more in the [Reduce memory usage](../optimization/memory) and [Torch 2.0](../optimization/torch2.0) guides.

docs/source/en/using-diffusers/marigold_usage.md

@@ -288,7 +288,7 @@ Speeding them up can be achieved by using a more efficient attention processor:
   depth = pipe(image, num_inference_steps=1)
 ```
 
-Finally, as suggested in [Optimizations](../optimization/fp16#torchcompile), enabling `torch.compile` can further enhance performance depending on 
+Finally, as suggested in [Optimizations](../optimization/torch2.0#torch.compile), enabling `torch.compile` can further enhance performance depending on 
 the target hardware.
 However, compilation incurs a significant overhead during the first pipeline invocation, making it beneficial only when 
 the same pipeline instance is called repeatedly, such as within a loop.

docs/source/en/using-diffusers/svd.md

@@ -63,7 +63,7 @@ export_to_video(frames, "generated.mp4", fps=7)
 
 ## torch.compile
 
-You can gain a 20-25% speedup at the expense of slightly increased memory by [compiling](../optimization/fp16#torchcompile) the UNet.
+You can gain a 20-25% speedup at the expense of slightly increased memory by [compiling](../optimization/torch2.0#torchcompile) the UNet.
 
 ```diff
 - pipe.enable_model_cpu_offload()

docs/source/en/using-diffusers/text-img2vid.md

@@ -1,4 +1,4 @@
-<!--Copyright 2024 The HuggingFace Team. All rights reserved.
+ <!--Copyright 2024 The HuggingFace Team. All rights reserved.
 
 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
@@ -12,436 +12,551 @@ specific language governing permissions and limitations under the License.
 
 # Video generation
 
-Video generation models extend image generation (can be considered a 1-frame video) to also process data related to space and time. Making sure all this data - text, space, time - remain consistent and aligned from frame-to-frame is a big challenge in generating long and high-resolution videos.
+Video generation models include a temporal dimension to bring images, or frames, together to create a video. These models are trained on large-scale datasets of high-quality text-video pairs to learn how to combine the modalities to ensure the generated video is coherent and realistic.
 
-Modern video models tackle this challenge with the diffusion transformer (DiT) architecture. This reduces computational costs and allows more efficient scaling to larger and higher-quality image and video data.
+[Explore](https://huggingface.co/models?other=video-generation) some of the more popular open-source video generation models available from Diffusers below.
 
-Check out what some of these video models are capable of below.
+<hfoptions id="popular-models">
+<hfoption id="CogVideoX">
 
-<hfoptions id="popular models">
-<hfoption id="Wan2.1">
+[CogVideoX](https://huggingface.co/collections/THUDM/cogvideo-66c08e62f1685a3ade464cce) uses a 3D causal Variational Autoencoder (VAE) to compress videos along the spatial and temporal dimensions, and it includes a stack of expert transformer blocks with a 3D full attention mechanism to better capture visual, semantic, and motion information in the data.
+
+The CogVideoX family also includes models capable of generating videos from images and videos in addition to text. The image-to-video models are indicated by **I2V** in the checkpoint name, and they should be used with the [`CogVideoXImageToVideoPipeline`]. The regular checkpoints support video-to-video through the [`CogVideoXVideoToVideoPipeline`].
+
+The example below demonstrates how to generate a video from an image and text prompt with [THUDM/CogVideoX-5b-I2V](https://huggingface.co/THUDM/CogVideoX-5b-I2V).
 
 ```py
-# pip install ftfy
 import torch
-import numpy as np
-from diffusers import AutoModel, WanPipeline
-from diffusers.hooks.group_offloading import apply_group_offloading
+from diffusers import CogVideoXImageToVideoPipeline
 from diffusers.utils import export_to_video, load_image
-from transformers import UMT5EncoderModel
 
-text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
-vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
-transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
-
-# group-offloading
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-apply_group_offloading(text_encoder,
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="block_level",
-    num_blocks_per_group=4
-)
-transformer.enable_group_offload(
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="leaf_level",
-    use_stream=True
-)
-
-pipeline = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-    vae=vae,
-    transformer=transformer,
-    text_encoder=text_encoder,
+prompt = "A vast, shimmering ocean flows gracefully under a twilight sky, its waves undulating in a mesmerizing dance of blues and greens. The surface glints with the last rays of the setting sun, casting golden highlights that ripple across the water. Seagulls soar above, their cries blending with the gentle roar of the waves. The horizon stretches infinitely, where the ocean meets the sky in a seamless blend of hues. Close-ups reveal the intricate patterns of the waves, capturing the fluidity and dynamic beauty of the sea in motion."
+image = load_image(image="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_rocket.png")
+pipe = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V",
     torch_dtype=torch.bfloat16
 )
-pipeline.to("cuda")
 
-prompt = """
-The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
-shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
-negative_prompt = """
-Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, 
-low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, 
-misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
-"""
+# reduce memory requirements 
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
 
-output = pipeline(
+video = pipe(
     prompt=prompt,
-    negative_prompt=negative_prompt,
-    num_frames=81,
-    guidance_scale=5.0,
+    image=image,
+    num_videos_per_prompt=1,
+    num_inference_steps=50,
+    num_frames=49,
+    guidance_scale=6,
+    generator=torch.Generator(device="cuda").manual_seed(42),
 ).frames[0]
-export_to_video(output, "output.mp4", fps=16)
+export_to_video(video, "output.mp4", fps=8)
 ```
 
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_rocket.png"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">initial image</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cogvideox/cogvideox_outrocket.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">generated video</figcaption>
+  </div>
+</div>
+
 </hfoption>
 <hfoption id="HunyuanVideo">
 
+> [!TIP]
+> HunyuanVideo is a 13B parameter model and requires a lot of memory. Refer to the HunyuanVideo [Quantization](../api/pipelines/hunyuan_video#quantization) guide to learn how to quantize the model. CogVideoX and LTX-Video are more lightweight options that can still generate high-quality videos.
+
+[HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo) features a dual-stream to single-stream diffusion transformer (DiT) for learning video and text tokens separately, and then subsequently concatenating the video and text tokens to combine their information. A single multimodal large language model (MLLM) serves as the text encoder, and videos are also spatio-temporally compressed with a 3D causal VAE.
+
 ```py
 import torch
-from diffusers importAutoModel, HunyuanVideoPipeline
-from diffusers.quantizers import PipelineQuantizationConfig
+from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
 from diffusers.utils import export_to_video
 
-# quantize weights to int4 with bitsandbytes
-pipeline_quant_config = PipelineQuantizationConfig(
-  quant_backend="bitsandbytes_4bit",
-  quant_kwargs={
-    "load_in_4bit": True,
-    "bnb_4bit_quant_type": "nf4",
-    "bnb_4bit_compute_dtype": torch.bfloat16
-    },
-  components_to_quantize=["transformer"]
+transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+    "hunyuanvideo-community/HunyuanVideo", subfolder="transformer", torch_dtype=torch.bfloat16
+)
+pipe = HunyuanVideoPipeline.from_pretrained(
+  "hunyuanvideo-community/HunyuanVideo", transformer=transformer, torch_dtype=torch.float16
 )
 
-pipeline = HunyuanVideoPipeline.from_pretrained(
-    "hunyuanvideo-community/HunyuanVideo",
-    quantization_config=pipeline_quant_config,
-    torch_dtype=torch.bfloat16,
-)
+# reduce memory requirements
+pipe.vae.enable_tiling()
+pipe.to("cuda")
 
-# model-offloading and tiling
-pipeline.enable_model_cpu_offload()
-pipeline.vae.enable_tiling()
-
-prompt = "A fluffy teddy bear sits on a bed of soft pillows surrounded by children's toys."
-video = pipeline(prompt=prompt, num_frames=61, num_inference_steps=30).frames[0]
+video = pipe(
+    prompt="A cat walks on the grass, realistic",
+    height=320,
+    width=512,
+    num_frames=61,
+    num_inference_steps=30,
+).frames[0]
 export_to_video(video, "output.mp4", fps=15)
 ```
 
+<div class="flex justify-center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/hunyuan-video-output.gif"/>
+</div>
+
 </hfoption>
 <hfoption id="LTX-Video">
 
-```py
-import torch
-from diffusers import LTXPipeline, AutoModel
-from diffusers.hooks import apply_group_offloading
-from diffusers.utils import export_to_video
-
-# fp8 layerwise weight-casting
-transformer = AutoModel.from_pretrained(
-    "Lightricks/LTX-Video",
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16
-)
-transformer.enable_layerwise_casting(
-    storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16
-)
-
-pipeline = LTXPipeline.from_pretrained("Lightricks/LTX-Video", transformer=transformer, torch_dtype=torch.bfloat16)
-
-# group-offloading
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
-apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
-apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
-
-prompt = """
-A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage
-"""
-negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"
-
-video = pipeline(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    width=768,
-    height=512,
-    num_frames=161,
-    decode_timestep=0.03,
-    decode_noise_scale=0.025,
-    num_inference_steps=50,
-).frames[0]
-export_to_video(video, "output.mp4", fps=24)
-```
-
-</hfoption>
-</hfoptions>
-
-This guide will cover video generation basics such as which parameters to configure and how to reduce their memory usage.
-
-> [!TIP]
-> If you're interested in learning more about how to use a specific model, please refer to their pipeline API model card.
-
-## Pipeline parameters
-
-There are several parameters to configure in the pipeline that'll affect video generation quality or speed. Experimenting with different parameter values is important for discovering the appropriate quality and speed tradeoff.
-
-### num_frames
-
-A frame is a still image that is played in a sequence of other frames to create motion or a video. Control the number of frames generated per second with `num_frames`. Increasing `num_frames` increases perceived motion smoothness and visual coherence, making it especially important for videos with dynamic content. A higher `num_frames` value also increases video duration.
-
-Some video models require more specific `num_frames` values for inference. For example, [`HunyuanVideoPipeline`] recommends calculating the `num_frames` with `(4 * num_frames) +1`. Always check a pipelines API model card to see if there is a recommended value.
+[LTX-Video (LTXV)](https://huggingface.co/Lightricks/LTX-Video) is a diffusion transformer (DiT) with a focus on speed. It generates 768x512 resolution videos at 24 frames per second (fps), enabling near real-time generation of high-quality videos. LTXV is relatively lightweight compared to other modern video generation models, making it possible to run on consumer GPUs.
 
 ```py
 import torch
 from diffusers import LTXPipeline
 from diffusers.utils import export_to_video
 
-pipeline = LTXPipeline.from_pretrained(
-    "Lightricks/LTX-Video", torch_dtype=torch.bfloat16
-).to("cuda")
+pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16).to("cuda")
 
-prompt = """
-A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman 
-with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The 
-camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and 
-natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be 
-real-life footage
-"""
-
-video = pipeline(
+prompt = "A man walks towards a window, looks out, and then turns around. He has short, dark hair, dark skin, and is wearing a brown coat over a red and gray scarf. He walks from left to right towards a window, his gaze fixed on something outside. The camera follows him from behind at a medium distance. The room is brightly lit, with white walls and a large window covered by a white curtain. As he approaches the window, he turns his head slightly to the left, then back to the right. He then turns his entire body to the right, facing the window. The camera remains stationary as he stands in front of the window. The scene is captured in real-life footage."
+video = pipe(
     prompt=prompt,
-    negative_prompt=negative_prompt,
-    width=768,
-    height=512,
+    width=704,
+    height=480,
     num_frames=161,
-    decode_timestep=0.03,
-    decode_noise_scale=0.025,
     num_inference_steps=50,
 ).frames[0]
 export_to_video(video, "output.mp4", fps=24)
 ```
 
-### guidance_scale
+<div class="flex justify-center">
+  <img src="https://huggingface.co/Lightricks/LTX-Video/resolve/main/media/ltx-video_example_00014.gif"/>
+</div>
 
-Guidance scale or "cfg" controls how closely the generated frames adhere to the input conditioning (text, image or both). Increasing `guidance_scale` generates frames that resemble the input conditions more closely and includes finer details, but risk introducing artifacts and reducing output diversity. Lower `guidance_scale` values encourages looser prompt adherence and increased output variety, but details may not be as great. If it's too low, it may ignore your prompt entirely and generate random noise.
-
-```py
-import torch
-from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
-from diffusers.utils import export_to_video
-
-pipeline = CogVideoXPipeline.from_pretrained(
-  "THUDM/CogVideoX-2b",
-  torch_dtype=torch.float16
-).to("cuda")
-
-prompt = """
-A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over
-a plush, blue carpet that mimics the waves of the sea. The ship's hull is painted a rich brown, 
-with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an 
-oceanic expanse. Surrounding the ship are various other toys and children's items, hinting at 
-a playful environment. The scene captures the innocence and imagination of childhood, 
-with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
-"""
-
-video = pipeline(
-  prompt=prompt,
-  guidance_scale=6,
-  num_inference_steps=50
-).frames[0]
-export_to_video(video, "output.mp4", fps=8)
-```
-
-### negative_prompt
-
-A negative prompt is useful for excluding things you don't want to see in the generated video. It is commonly used to refine the quality and alignment of the generated video by pushing the model away from undesirable elements like "blurry, distorted, ugly". This can create cleaner and more focused videos.
-
-```py
-# pip install ftfy
-import torch
-from diffusers import WanPipeline
-from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-from diffusers.utils import export_to_video
-
-vae = AutoencoderKLWan.from_pretrained(
-  "Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
-)
-pipeline = WanPipeline.from_pretrained(
-  "Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, torch_dtype=torch.bfloat16
-)
-pipeline.scheduler = UniPCMultistepScheduler.from_config(
-  pipeline.scheduler.config, flow_shift=5.0
-)
-pipeline.to("cuda")
-
-pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
-pipeline.set_adapters("steamboat-willie")
-
-pipeline.enable_model_cpu_offload()
-
-# use "steamboat willie style" to trigger the LoRA
-prompt = """
-steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts 
-dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
-
-output = pipeline(
-  prompt=prompt,
-  num_frames=81,
-  guidance_scale=5.0,
-).frames[0]
-export_to_video(output, "output.mp4", fps=16)
-```
-
-## Reduce memory usage
-
-Recent video models like [`HunyuanVideoPipeline`] and [`WanPipeline`], which have 10B+ parameters, require a lot of memory and it often exceeds the memory availabe on consumer hardware. Diffusers offers several techniques for reducing the memory requirements of these large models.
+</hfoption>
+<hfoption id="Mochi-1">
 
 > [!TIP]
-> Refer to the [Reduce memory usage](../optimization/memory) guide for more details about other memory saving techniques.
+> Mochi-1 is a 10B parameter model and requires a lot of memory. Refer to the Mochi [Quantization](../api/pipelines/mochi#quantization) guide to learn how to quantize the model. CogVideoX and LTX-Video are more lightweight options that can still generate high-quality videos.
 
-One of these techniques is [group-offloading](../optimization/memory#group-offloading), which offloads groups of internal model layers (such as `torch.nn.Sequential`) to the CPU when it isn't being used. These layers are only loaded when they're needed for computation to avoid storing **all** the model components on the GPU. For a 14B parameter model like [`WanPipeline`], group-offloading can lower the required memory to ~13GB of VRAM.
+[Mochi-1](https://huggingface.co/genmo/mochi-1-preview) introduces the Asymmetric Diffusion Transformer (AsymmDiT) and Asymmetric Variational Autoencoder (AsymmVAE) to reduces memory requirements. AsymmVAE causally compresses videos 128x to improve memory efficiency, and AsymmDiT jointly attends to the compressed video tokens and user text tokens. This model is noted for generating videos with high-quality motion dynamics and strong prompt adherence.
 
 ```py
-# pip install ftfy
 import torch
-import numpy as np
-from diffusers import AutoModel, WanPipeline
-from diffusers.hooks.group_offloading import apply_group_offloading
-from diffusers.utils import export_to_video, load_image
-from transformers import UMT5EncoderModel
-
-text_encoder = UMT5EncoderModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", torch_dtype=torch.bfloat16)
-vae = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32)
-transformer = AutoModel.from_pretrained("Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16)
-
-# group-offloading
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-apply_group_offloading(text_encoder,
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="block_level",
-    num_blocks_per_group=4
-)
-transformer.enable_group_offload(
-    onload_device=onload_device,
-    offload_device=offload_device,
-    offload_type="leaf_level",
-    use_stream=True
-)
-
-pipeline = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers",
-    vae=vae,
-    transformer=transformer,
-    text_encoder=text_encoder,
-    torch_dtype=torch.bfloat16
-)
-pipeline.to("cuda")
-
-prompt = """
-The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic 
-shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
-negative_prompt = """
-Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, 
-low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, 
-misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
-"""
-
-output = pipeline(
-    prompt=prompt,
-    negative_prompt=negative_prompt,
-    num_frames=81,
-    guidance_scale=5.0,
-).frames[0]
-export_to_video(output, "output.mp4", fps=16)
-```
-
-Another option for reducing memory is to consider quantizing a model, which stores the model weights in a lower precision data type. However, quantization may impact video quality depending on the specific video model. Refer to the quantization [Overivew](../quantization/overview) to learn more about the different supported quantization backends.
-
-The example below uses [bitsandbytes](../quantization/bitsandbytes) to quantize a model.
-
-```py
-# pip install ftfy
-
-import torch
-from diffusers import WanPipeline
-from diffusers import AutoModel, WanPipeline
-from diffusers.quantizers import PipelineQuantizationConfig
-from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-from transformers import UMT5EncoderModel
+from diffusers import MochiPipeline
 from diffusers.utils import export_to_video
 
-# quantize transformer and text encoder weights with bitsandbytes
-pipeline_quant_config = PipelineQuantizationConfig(
-  quant_backend="bitsandbytes_4bit",
-  quant_kwargs={"load_in_4bit": True},
-  components_to_quantize=["transformer", "text_encoder"]
+pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", variant="bf16", torch_dtype=torch.bfloat16)
+
+# reduce memory requirements
+pipe.enable_model_cpu_offload()
+pipe.enable_vae_tiling()
+
+prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
+video = pipe(prompt, num_frames=84).frames[0]
+export_to_video(video, "output.mp4", fps=30)
+```
+
+<div class="flex justify-center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/mochi-video-output.gif"/>
+</div>
+
+</hfoption>
+<hfoption id="StableVideoDiffusion">
+
+[StableVideoDiffusion (SVD)](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt) is based on the Stable Diffusion 2.1 model and it is trained on images, then low-resolution videos, and finally a smaller dataset of high-resolution videos. This model generates a short 2-4 second video from an initial image.
+
+```py
+import torch
+from diffusers import StableVideoDiffusionPipeline
+from diffusers.utils import load_image, export_to_video
+
+pipeline = StableVideoDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
 )
 
-vae = AutoModel.from_pretrained(
-  "Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="vae", torch_dtype=torch.float32
-)
-pipeline = WanPipeline.from_pretrained(
-  "Wan-AI/Wan2.1-T2V-14B-Diffusers", vae=vae, quantization_config=pipeline_quant_config, torch_dtype=torch.bfloat16
-)
-pipeline.scheduler = UniPCMultistepScheduler.from_config(
-  pipeline.scheduler.config, flow_shift=5.0
-)
-pipeline.to("cuda")
-
-pipeline.load_lora_weights("benjamin-paine/steamboat-willie-14b", adapter_name="steamboat-willie")
-pipeline.set_adapters("steamboat-willie")
-
+# reduce memory requirements
 pipeline.enable_model_cpu_offload()
 
-# use "steamboat willie style" to trigger the LoRA
-prompt = """
-steamboat willie style, golden era animation, The camera rushes from far to near in a low-angle shot, 
-revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in 
-for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground. 
-Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts 
-dynamic shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
-"""
+image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png")
+image = image.resize((1024, 576))
 
-output = pipeline(
-  prompt=prompt,
-  num_frames=81,
-  guidance_scale=5.0,
-).frames[0]
-export_to_video(output, "output.mp4", fps=16)
+generator = torch.manual_seed(42)
+frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
+export_to_video(frames, "generated.mp4", fps=7)
 ```
 
-## Inference speed
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">initial image</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/output_rocket.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">generated video</figcaption>
+  </div>
+</div>
 
-[torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial_.html) can speedup inference by using optimized kernels. Compilation takes longer the first time, but once compiled, it is much faster. It is best to compile the pipeline once, and then use the pipeline multiple times without changing anything. A change, such as in the image size, triggers recompilation.
+</hfoption>
+<hfoption id="AnimateDiff">
 
-The example below compiles the transformer in the pipeline and uses the `"max-autotune"` mode to maximize performance.
+[AnimateDiff](https://huggingface.co/guoyww/animatediff) is an adapter model that inserts a motion module into a pretrained diffusion model to animate an image. The adapter is trained on video clips to learn motion which is used to condition the generation process to create a video. It is faster and easier to only train the adapter and it can be loaded into most diffusion models, effectively turning them into “video models”.
+
+Load a `MotionAdapter` and pass it to the [`AnimateDiffPipeline`].
 
 ```py
 import torch
-from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
-from diffusers.utils import export_to_video
+from diffusers import AnimateDiffPipeline, DDIMScheduler, MotionAdapter
+from diffusers.utils import export_to_gif
 
-pipeline = CogVideoXPipeline.from_pretrained(
-  "THUDM/CogVideoX-2b",
-  torch_dtype=torch.float16
+adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)
+pipeline = AnimateDiffPipeline.from_pretrained("emilianJR/epiCRealism", motion_adapter=adapter, torch_dtype=torch.float16)
+scheduler = DDIMScheduler.from_pretrained(
+    "emilianJR/epiCRealism",
+    subfolder="scheduler",
+    clip_sample=False,
+    timestep_spacing="linspace",
+    beta_schedule="linear",
+    steps_offset=1,
+)
+pipeline.scheduler = scheduler
+
+# reduce memory requirements
+pipeline.enable_vae_slicing()
+pipeline.enable_model_cpu_offload()
+
+output = pipeline(
+    prompt="A space rocket with trails of smoke behind it launching into space from the desert, 4k, high resolution",
+    negative_prompt="bad quality, worse quality, low resolution",
+    num_frames=16,
+    guidance_scale=7.5,
+    num_inference_steps=50,
+    generator=torch.Generator("cpu").manual_seed(49),
+)
+frames = output.frames[0]
+export_to_gif(frames, "animation.gif")
+```
+
+<div class="flex justify-center">
+  <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff.gif"/>
+</div>
+
+</hfoption>
+</hfoptions>
+
+## Configure model parameters
+
+There are a few important parameters you can configure in the pipeline that'll affect the video generation process and quality. Let's take a closer look at what these parameters do and how changing them affects the output.
+
+### Number of frames
+
+The `num_frames` parameter determines how many video frames are generated per second. A frame is an image that is played in a sequence of other frames to create motion or a video. This affects video length because the pipeline generates a certain number of frames per second (check a pipeline's API reference for the default value). To increase the video duration, you'll need to increase the `num_frames` parameter.
+
+```py
+import torch
+from diffusers import StableVideoDiffusionPipeline
+from diffusers.utils import load_image, export_to_video
+
+pipeline = StableVideoDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
+)
+pipeline.enable_model_cpu_offload()
+
+image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/svd/rocket.png")
+image = image.resize((1024, 576))
+
+generator = torch.manual_seed(42)
+frames = pipeline(image, decode_chunk_size=8, generator=generator, num_frames=25).frames[0]
+export_to_video(frames, "generated.mp4", fps=7)
+```
+
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/num_frames_14.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">num_frames=14</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/num_frames_25.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">num_frames=25</figcaption>
+  </div>
+</div>
+
+### Guidance scale
+
+The `guidance_scale` parameter controls how closely aligned the generated video and text prompt or initial image is. A higher `guidance_scale` value means your generated video is more aligned with the text prompt or initial image, while a lower `guidance_scale` value means your generated video is less aligned which could give the model more "creativity" to interpret the conditioning input.
+
+<Tip>
+
+SVD uses the `min_guidance_scale` and `max_guidance_scale` parameters for applying guidance to the first and last frames respectively.
+
+</Tip>
+
+```py
+import torch
+from diffusers import I2VGenXLPipeline
+from diffusers.utils import export_to_gif, load_image
+
+pipeline = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")
+pipeline.enable_model_cpu_offload()
+
+image_url = "https://huggingface.co/datasets/diffusers/docs-images/resolve/main/i2vgen_xl_images/img_0009.png"
+image = load_image(image_url).convert("RGB")
+
+prompt = "Papers were floating in the air on a table in the library"
+negative_prompt = "Distorted, discontinuous, Ugly, blurry, low resolution, motionless, static, disfigured, disconnected limbs, Ugly faces, incomplete arms"
+generator = torch.manual_seed(0)
+
+frames = pipeline(
+    prompt=prompt,
+    image=image,
+    num_inference_steps=50,
+    negative_prompt=negative_prompt,
+    guidance_scale=1.0,
+    generator=generator
+).frames[0]
+export_to_gif(frames, "i2v.gif")
+```
+
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/i2vgen-xl-example.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">guidance_scale=9.0</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guidance_scale_1.0.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">guidance_scale=1.0</figcaption>
+  </div>
+</div>
+
+### Negative prompt
+
+A negative prompt deters the model from generating things you don’t want it to. This parameter is commonly used to improve overall generation quality by removing poor or bad features such as “low resolution” or “bad details”.
+
+```py
+import torch
+from diffusers import AnimateDiffPipeline, DDIMScheduler, MotionAdapter
+from diffusers.utils import export_to_gif
+
+adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)
+
+pipeline = AnimateDiffPipeline.from_pretrained("emilianJR/epiCRealism", motion_adapter=adapter, torch_dtype=torch.float16)
+scheduler = DDIMScheduler.from_pretrained(
+    "emilianJR/epiCRealism",
+    subfolder="scheduler",
+    clip_sample=False,
+    timestep_spacing="linspace",
+    beta_schedule="linear",
+    steps_offset=1,
+)
+pipeline.scheduler = scheduler
+pipeline.enable_vae_slicing()
+pipeline.enable_model_cpu_offload()
+
+output = pipeline(
+    prompt="360 camera shot of a sushi roll in a restaurant",
+    negative_prompt="Distorted, discontinuous, ugly, blurry, low resolution, motionless, static",
+    num_frames=16,
+    guidance_scale=7.5,
+    num_inference_steps=50,
+    generator=torch.Generator("cpu").manual_seed(0),
+)
+frames = output.frames[0]
+export_to_gif(frames, "animation.gif")
+```
+
+<div class="flex gap-4">
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff_no_neg.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">no negative prompt</figcaption>
+  </div>
+  <div>
+    <img class="rounded-xl" src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff_neg.gif"/>
+    <figcaption class="mt-2 text-center text-sm text-gray-500">negative prompt applied</figcaption>
+  </div>
+</div>
+
+### Model-specific parameters
+
+There are some pipeline parameters that are unique to each model such as adjusting the motion in a video or adding noise to the initial image.
+
+<hfoptions id="special-parameters">
+<hfoption id="Stable Video Diffusion">
+
+Stable Video Diffusion provides additional micro-conditioning for the frame rate with the `fps` parameter and for motion with the `motion_bucket_id` parameter. Together, these parameters allow for adjusting the amount of motion in the generated video.
+
+There is also a `noise_aug_strength` parameter that increases the amount of noise added to the initial image. Varying this parameter affects how similar the generated video and initial image are. A higher `noise_aug_strength` also increases the amount of motion. To learn more, read the [Micro-conditioning](../using-diffusers/svd#micro-conditioning) guide.
+
+</hfoption>
+<hfoption id="Text2Video-Zero">
+
+Text2Video-Zero computes the amount of motion to apply to each frame from randomly sampled latents. You can use the `motion_field_strength_x` and `motion_field_strength_y` parameters to control the amount of motion to apply to the x and y-axes of the video. The parameters `t0` and `t1` are the timesteps to apply motion to the latents.
+
+</hfoption>
+</hfoptions>
+
+## Control video generation
+
+Video generation can be controlled similar to how text-to-image, image-to-image, and inpainting can be controlled with a [`ControlNetModel`]. The only difference is you need to use the [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor`] so each frame attends to the first frame.
+
+### Text2Video-Zero
+
+Text2Video-Zero video generation can be conditioned on pose and edge images for even greater control over a subject's motion in the generated video or to preserve the identity of a subject/object in the video. You can also use Text2Video-Zero with [InstructPix2Pix](../api/pipelines/pix2pix) for editing videos with text.
+
+<hfoptions id="t2v-zero">
+<hfoption id="pose control">
+
+Start by downloading a video and extracting the pose images from it.
+
+```py
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import imageio
+
+filename = "__assets__/poses_skeleton_gifs/dance1_corr.mp4"
+repo_id = "PAIR/Text2Video-Zero"
+video_path = hf_hub_download(repo_type="space", repo_id=repo_id, filename=filename)
+
+reader = imageio.get_reader(video_path, "ffmpeg")
+frame_count = 8
+pose_images = [Image.fromarray(reader.get_data(i)) for i in range(frame_count)]
+```
+
+Load a [`ControlNetModel`] for pose estimation and a checkpoint into the [`StableDiffusionControlNetPipeline`]. Then you'll use the [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor`] for the UNet and ControlNet.
+
+```py
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero import CrossFrameAttnProcessor
+
+model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16)
+pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+    model_id, controlnet=controlnet, torch_dtype=torch.float16
 ).to("cuda")
 
-# torch.compile
-pipeline.transformer.to(memory_format=torch.channels_last)
-pipeline.transformer = torch.compile(
-    pipeline.transformer, mode="max-autotune", fullgraph=True
-)
+pipeline.unet.set_attn_processor(CrossFrameAttnProcessor(batch_size=2))
+pipeline.controlnet.set_attn_processor(CrossFrameAttnProcessor(batch_size=2))
+```
 
-prompt = """
-A detailed wooden toy ship with intricately carved masts and sails is seen gliding smoothly over a plush, blue carpet that mimics the waves of the sea. 
-The ship's hull is painted a rich brown, with tiny windows. The carpet, soft and textured, provides a perfect backdrop, resembling an oceanic expanse. 
-Surrounding the ship are various other toys and children's items, hinting at a playful environment. The scene captures the innocence and imagination of childhood, 
-with the toy ship's journey symbolizing endless adventures in a whimsical, indoor setting.
-"""
+Fix the latents for all the frames, and then pass your prompt and extracted pose images to the model to generate a video.
 
-video = pipeline(
-  prompt=prompt,
-  guidance_scale=6,
-  num_inference_steps=50
-).frames[0]
-export_to_video(video, "output.mp4", fps=8)
-```
+```py
+latents = torch.randn((1, 4, 64, 64), device="cuda", dtype=torch.float16).repeat(len(pose_images), 1, 1, 1)
+
+prompt = "Darth Vader dancing in a desert"
+result = pipeline(prompt=[prompt] * len(pose_images), image=pose_images, latents=latents).images
+imageio.mimsave("video.mp4", result, fps=4)
+```
+
+</hfoption>
+<hfoption id="edge control">
+
+Download a video and extract the edges from it.
+
+```py
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import imageio
+
+filename = "__assets__/poses_skeleton_gifs/dance1_corr.mp4"
+repo_id = "PAIR/Text2Video-Zero"
+video_path = hf_hub_download(repo_type="space", repo_id=repo_id, filename=filename)
+
+reader = imageio.get_reader(video_path, "ffmpeg")
+frame_count = 8
+pose_images = [Image.fromarray(reader.get_data(i)) for i in range(frame_count)]
+```
+
+Load a [`ControlNetModel`] for canny edge and a checkpoint into the [`StableDiffusionControlNetPipeline`]. Then you'll use the [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor`] for the UNet and ControlNet.
+
+```py
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero import CrossFrameAttnProcessor
+
+model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
+controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
+pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+    model_id, controlnet=controlnet, torch_dtype=torch.float16
+).to("cuda")
+
+pipeline.unet.set_attn_processor(CrossFrameAttnProcessor(batch_size=2))
+pipeline.controlnet.set_attn_processor(CrossFrameAttnProcessor(batch_size=2))
+```
+
+Fix the latents for all the frames, and then pass your prompt and extracted edge images to the model to generate a video.
+
+```py
+latents = torch.randn((1, 4, 64, 64), device="cuda", dtype=torch.float16).repeat(len(pose_images), 1, 1, 1)
+
+prompt = "Darth Vader dancing in a desert"
+result = pipeline(prompt=[prompt] * len(pose_images), image=pose_images, latents=latents).images
+imageio.mimsave("video.mp4", result, fps=4)
+```
+
+</hfoption>
+<hfoption id="InstructPix2Pix">
+
+InstructPix2Pix allows you to use text to describe the changes you want to make to the video. Start by downloading and reading a video.
+
+```py
+from huggingface_hub import hf_hub_download
+from PIL import Image
+import imageio
+
+filename = "__assets__/pix2pix video/camel.mp4"
+repo_id = "PAIR/Text2Video-Zero"
+video_path = hf_hub_download(repo_type="space", repo_id=repo_id, filename=filename)
+
+reader = imageio.get_reader(video_path, "ffmpeg")
+frame_count = 8
+video = [Image.fromarray(reader.get_data(i)) for i in range(frame_count)]
+```
+
+Load the [`StableDiffusionInstructPix2PixPipeline`] and set the [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.CrossFrameAttnProcessor`] for the UNet.
+
+```py
+import torch
+from diffusers import StableDiffusionInstructPix2PixPipeline
+from diffusers.pipelines.text_to_video_synthesis.pipeline_text_to_video_zero import CrossFrameAttnProcessor
+
+pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", torch_dtype=torch.float16).to("cuda")
+pipeline.unet.set_attn_processor(CrossFrameAttnProcessor(batch_size=3))
+```
+
+Pass a prompt describing the change you want to apply to the video.
+
+```py
+prompt = "make it Van Gogh Starry Night style"
+result = pipeline(prompt=[prompt] * len(video), image=video).images
+imageio.mimsave("edited_video.mp4", result, fps=4)
+```
+
+</hfoption>
+</hfoptions>
+
+## Optimize
+
+Video generation requires a lot of memory because you're generating many video frames at once. You can reduce your memory requirements at the expense of some inference speed. Try:
+
+1. offloading pipeline components that are no longer needed to the CPU
+2. feed-forward chunking runs the feed-forward layer in a loop instead of all at once
+3. break up the number of frames the VAE has to decode into chunks instead of decoding them all at once
+
+```diff
+- pipeline.enable_model_cpu_offload()
+- frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]
++ pipeline.enable_model_cpu_offload()
++ pipeline.unet.enable_forward_chunking()
++ frames = pipeline(image, decode_chunk_size=2, generator=generator, num_frames=25).frames[0]
+```
+
+If memory is not an issue and you want to optimize for speed, try wrapping the UNet with [`torch.compile`](../optimization/torch2.0#torchcompile).
+
+```diff
+- pipeline.enable_model_cpu_offload()
++ pipeline.to("cuda")
++ pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
+```
+
+## Quantization
+
+Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model.
+
+Refer to the [Quantization](../../quantization/overview) to learn more about supported quantization backends (bitsandbytes, torchao, gguf) and selecting a quantization backend that supports your use case.

examples/advanced_diffusion_training/README.md

@@ -128,7 +128,6 @@ You can also load a dataset straight from by specifying it's name in `dataset_na
 Look [here](https://huggingface.co/blog/sdxl_lora_advanced_script#custom-captioning) for more info on creating/loading your own caption dataset.
 
 - **optimizer**: for this example, we'll use [prodigy](https://huggingface.co/blog/sdxl_lora_advanced_script#adaptive-optimizers) - an adaptive optimizer
-  - To use Prodigy, please make sure to install the prodigyopt library: `pip install prodigyopt`
 - **pivotal tuning**
 - **min SNR gamma**
 

examples/advanced_diffusion_training/README_flux.md

@@ -143,8 +143,7 @@ Now we'll simply specify the name of the dataset and caption column (in this cas
 You can also load a dataset straight from by specifying it's name in `dataset_name`.
 Look [here](https://huggingface.co/blog/sdxl_lora_advanced_script#custom-captioning) for more info on creating/loading your own caption dataset.
 
-- **optimizer**: for this example, we'll use [prodigy](https://huggingface.co/blog/sdxl_lora_advanced_script#adaptive-optimizers) - an adaptive optimizer 
-    - To use Prodigy, please make sure to install the prodigyopt library: `pip install prodigyopt`
+- **optimizer**: for this example, we'll use [prodigy](https://huggingface.co/blog/sdxl_lora_advanced_script#adaptive-optimizers) - an adaptive optimizer
 - **pivotal tuning**
 
 ### Example #1: Pivotal tuning

examples/cogvideo/train_cogvideox_image_to_video_lora.py

@@ -555,7 +555,7 @@ class VideoDataset(Dataset):
 
         if any(not path.is_file() for path in instance_videos):
             raise ValueError(
-                "Expected '--video_column' to be a path to a file in `--instance_data_root` containing line-separated paths to video data but found at least one path that is not a valid file."
+                "Expected '--video_column' to be a path to a file in `--instance_data_root` containing line-separated paths to video data but found atleast one path that is not a valid file."
             )
 
         return instance_prompts, instance_videos

examples/cogvideo/train_cogvideox_lora.py

@@ -539,7 +539,7 @@ class VideoDataset(Dataset):
 
         if any(not path.is_file() for path in instance_videos):
             raise ValueError(
-                "Expected '--video_column' to be a path to a file in `--instance_data_root` containing line-separated paths to video data but found at least one path that is not a valid file."
+                "Expected '--video_column' to be a path to a file in `--instance_data_root` containing line-separated paths to video data but found atleast one path that is not a valid file."
             )
 
         return instance_prompts, instance_videos

examples/controlnet/train_controlnet.py

@@ -178,11 +178,11 @@ def log_validation(
         else:
             logger.warning(f"image logging not implemented for {tracker.name}")
 
-    del pipeline
-    gc.collect()
-    torch.cuda.empty_cache()
+        del pipeline
+        gc.collect()
+        torch.cuda.empty_cache()
 
-    return image_logs
+        return image_logs
 
 
 def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str, revision: str):

examples/controlnet/train_controlnet_flux.py

@@ -192,9 +192,9 @@ def log_validation(
         else:
             logger.warning(f"image logging not implemented for {tracker.name}")
 
-    del pipeline
-    free_memory()
-    return image_logs
+        del pipeline
+        free_memory()
+        return image_logs
 
 
 def save_model_card(repo_id: str, image_logs=None, base_model=str, repo_folder=None):

examples/controlnet/train_controlnet_sd3.py

@@ -199,13 +199,13 @@ def log_validation(controlnet, args, accelerator, weight_dtype, step, is_final_v
         else:
             logger.warning(f"image logging not implemented for {tracker.name}")
 
-    del pipeline
-    free_memory()
+        del pipeline
+        free_memory()
 
-    if not is_final_validation:
-        controlnet.to(accelerator.device)
+        if not is_final_validation:
+            controlnet.to(accelerator.device)
 
-    return image_logs
+        return image_logs
 
 
 # Copied from dreambooth sd3 example

examples/controlnet/train_controlnet_sdxl.py

@@ -201,11 +201,11 @@ def log_validation(vae, unet, controlnet, args, accelerator, weight_dtype, step,
         else:
             logger.warning(f"image logging not implemented for {tracker.name}")
 
-    del pipeline
-    gc.collect()
-    torch.cuda.empty_cache()
+        del pipeline
+        gc.collect()
+        torch.cuda.empty_cache()
 
-    return image_logs
+        return image_logs
 
 
 def import_model_class_from_model_name_or_path(

examples/dreambooth/README_flux.md

@@ -134,7 +134,7 @@ Note also that we use PEFT library as backend for LoRA training, make sure to ha
 Prodigy is an adaptive optimizer that dynamically adjusts the learning rate learned parameters based on past gradients, allowing for more efficient convergence. 
 By using prodigy we can "eliminate" the need for manual learning rate tuning. read more [here](https://huggingface.co/blog/sdxl_lora_advanced_script#adaptive-optimizers).
 
-to use prodigy, first make sure to install the prodigyopt library: `pip install prodigyopt`, and then specify -
+to use prodigy, specify
 ```bash
 --optimizer="prodigy"
 ```

examples/flux-control/train_control_flux.py

@@ -915,7 +915,7 @@ def main(args):
         args.lr_scheduler,
         optimizer=optimizer,
         num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
-        num_training_steps=num_training_steps_for_scheduler,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
         num_cycles=args.lr_num_cycles,
         power=args.lr_power,
     )

examples/flux-control/train_control_lora_flux.py

@@ -1060,7 +1060,7 @@ def main(args):
         args.lr_scheduler,
         optimizer=optimizer,
         num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
-        num_training_steps=num_training_steps_for_scheduler,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
         num_cycles=args.lr_num_cycles,
         power=args.lr_power,
     )

examples/research_projects/multi_subject_dreambooth_inpainting/README.md

@@ -73,7 +73,7 @@ accelerate launch train_multi_subject_dreambooth_inpaint.py \
 
 ## 3. Results
 
-A [![Weights & Biases](https://img.shields.io/badge/Weights%20&%20Biases-Report-blue)](https://wandb.ai/gzguevara/uncategorized/reports/Multi-Subject-Dreambooth-for-Inpainting--Vmlldzo2MzY5NDQ4?accessToken=y0nya2d7baguhbryxaikbfr1203amvn1jsmyl07vk122mrs7tnph037u1nqgse8t) is provided showing the training progress by every 50 steps. Note, the reported weights & biases run was performed on a A100 GPU with the following stetting:
+A [![Weights & Biases](https://img.shields.io/badge/Weights%20&%20Biases-Report-blue)](https://wandb.ai/gzguevara/uncategorized/reports/Multi-Subject-Dreambooth-for-Inpainting--Vmlldzo2MzY5NDQ4?accessToken=y0nya2d7baguhbryxaikbfr1203amvn1jsmyl07vk122mrs7tnph037u1nqgse8t) is provided showing the training progress by every 50 steps. Note, the reported  weights & baises run was performed on a A100 GPU with the following stetting:
 
 ```bash
 accelerate launch train_multi_subject_dreambooth_inpaint.py \

examples/textual_inversion/textual_inversion_sdxl.py

@@ -793,22 +793,17 @@ def main():
     )
 
     # Scheduler and math around the number of training steps.
-    # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
-    num_warmup_steps_for_scheduler = args.lr_warmup_steps * accelerator.num_processes
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if args.max_train_steps is None:
-        len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
-        num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
-        num_training_steps_for_scheduler = (
-            args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
-        )
-    else:
-        num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
 
     lr_scheduler = get_scheduler(
         args.lr_scheduler,
         optimizer=optimizer,
-        num_warmup_steps=num_warmup_steps_for_scheduler,
-        num_training_steps=num_training_steps_for_scheduler,
+        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+        num_training_steps=args.max_train_steps * accelerator.num_processes,
         num_cycles=args.lr_num_cycles,
     )
 
@@ -834,14 +829,8 @@ def main():
 
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if args.max_train_steps is None:
+    if overrode_max_train_steps:
         args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-        if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
-            logger.warning(
-                f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
-                f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
-                f"This inconsistency may result in the learning rate scheduler not functioning properly."
-            )
     # Afterwards we recalculate our number of training epochs
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 

scripts/convert_cosmos_to_diffusers.py

@@ -7,16 +7,7 @@ from accelerate import init_empty_weights
 from huggingface_hub import snapshot_download
 from transformers import T5EncoderModel, T5TokenizerFast
 
-from diffusers import (
-    AutoencoderKLCosmos,
-    AutoencoderKLWan,
-    Cosmos2TextToImagePipeline,
-    Cosmos2VideoToWorldPipeline,
-    CosmosTextToWorldPipeline,
-    CosmosTransformer3DModel,
-    CosmosVideoToWorldPipeline,
-    EDMEulerScheduler,
-)
+from diffusers import AutoencoderKLCosmos, CosmosTextToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler
 
 
 def remove_keys_(key: str, state_dict: Dict[str, Any]):
@@ -38,7 +29,7 @@ def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]):
     state_dict[new_key] = state_dict.pop(key)
 
 
-TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 = {
+TRANSFORMER_KEYS_RENAME_DICT = {
     "t_embedder.1": "time_embed.t_embedder",
     "affline_norm": "time_embed.norm",
     ".blocks.0.block.attn": ".attn1",
@@ -65,7 +56,7 @@ TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 = {
     "final_layer.linear": "proj_out",
 }
 
-TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 = {
+TRANSFORMER_SPECIAL_KEYS_REMAP = {
     "blocks.block": rename_transformer_blocks_,
     "logvar.0.freqs": remove_keys_,
     "logvar.0.phases": remove_keys_,
@@ -73,45 +64,6 @@ TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 = {
     "pos_embedder.seq": remove_keys_,
 }
 
-TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 = {
-    "t_embedder.1": "time_embed.t_embedder",
-    "t_embedding_norm": "time_embed.norm",
-    "blocks": "transformer_blocks",
-    "adaln_modulation_self_attn.1": "norm1.linear_1",
-    "adaln_modulation_self_attn.2": "norm1.linear_2",
-    "adaln_modulation_cross_attn.1": "norm2.linear_1",
-    "adaln_modulation_cross_attn.2": "norm2.linear_2",
-    "adaln_modulation_mlp.1": "norm3.linear_1",
-    "adaln_modulation_mlp.2": "norm3.linear_2",
-    "self_attn": "attn1",
-    "cross_attn": "attn2",
-    "q_proj": "to_q",
-    "k_proj": "to_k",
-    "v_proj": "to_v",
-    "output_proj": "to_out.0",
-    "q_norm": "norm_q",
-    "k_norm": "norm_k",
-    "mlp.layer1": "ff.net.0.proj",
-    "mlp.layer2": "ff.net.2",
-    "x_embedder.proj.1": "patch_embed.proj",
-    # "extra_pos_embedder": "learnable_pos_embed",
-    "final_layer.adaln_modulation.1": "norm_out.linear_1",
-    "final_layer.adaln_modulation.2": "norm_out.linear_2",
-    "final_layer.linear": "proj_out",
-}
-
-TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 = {
-    "accum_video_sample_counter": remove_keys_,
-    "accum_image_sample_counter": remove_keys_,
-    "accum_iteration": remove_keys_,
-    "accum_train_in_hours": remove_keys_,
-    "pos_embedder.seq": remove_keys_,
-    "pos_embedder.dim_spatial_range": remove_keys_,
-    "pos_embedder.dim_temporal_range": remove_keys_,
-    "_extra_state": remove_keys_,
-}
-
-
 TRANSFORMER_CONFIGS = {
     "Cosmos-1.0-Diffusion-7B-Text2World": {
         "in_channels": 16,
@@ -173,66 +125,6 @@ TRANSFORMER_CONFIGS = {
         "concat_padding_mask": True,
         "extra_pos_embed_type": "learnable",
     },
-    "Cosmos-2.0-Diffusion-2B-Text2Image": {
-        "in_channels": 16,
-        "out_channels": 16,
-        "num_attention_heads": 16,
-        "attention_head_dim": 128,
-        "num_layers": 28,
-        "mlp_ratio": 4.0,
-        "text_embed_dim": 1024,
-        "adaln_lora_dim": 256,
-        "max_size": (128, 240, 240),
-        "patch_size": (1, 2, 2),
-        "rope_scale": (1.0, 4.0, 4.0),
-        "concat_padding_mask": True,
-        "extra_pos_embed_type": None,
-    },
-    "Cosmos-2.0-Diffusion-14B-Text2Image": {
-        "in_channels": 16,
-        "out_channels": 16,
-        "num_attention_heads": 40,
-        "attention_head_dim": 128,
-        "num_layers": 36,
-        "mlp_ratio": 4.0,
-        "text_embed_dim": 1024,
-        "adaln_lora_dim": 256,
-        "max_size": (128, 240, 240),
-        "patch_size": (1, 2, 2),
-        "rope_scale": (1.0, 4.0, 4.0),
-        "concat_padding_mask": True,
-        "extra_pos_embed_type": None,
-    },
-    "Cosmos-2.0-Diffusion-2B-Video2World": {
-        "in_channels": 16 + 1,
-        "out_channels": 16,
-        "num_attention_heads": 16,
-        "attention_head_dim": 128,
-        "num_layers": 28,
-        "mlp_ratio": 4.0,
-        "text_embed_dim": 1024,
-        "adaln_lora_dim": 256,
-        "max_size": (128, 240, 240),
-        "patch_size": (1, 2, 2),
-        "rope_scale": (1.0, 3.0, 3.0),
-        "concat_padding_mask": True,
-        "extra_pos_embed_type": None,
-    },
-    "Cosmos-2.0-Diffusion-14B-Video2World": {
-        "in_channels": 16 + 1,
-        "out_channels": 16,
-        "num_attention_heads": 40,
-        "attention_head_dim": 128,
-        "num_layers": 36,
-        "mlp_ratio": 4.0,
-        "text_embed_dim": 1024,
-        "adaln_lora_dim": 256,
-        "max_size": (128, 240, 240),
-        "patch_size": (1, 2, 2),
-        "rope_scale": (20 / 24, 2.0, 2.0),
-        "concat_padding_mask": True,
-        "extra_pos_embed_type": None,
-    },
 }
 
 VAE_KEYS_RENAME_DICT = {
@@ -324,18 +216,9 @@ def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
     return state_dict
 
 
-def convert_transformer(transformer_type: str, ckpt_path: str, weights_only: bool = True):
+def convert_transformer(transformer_type: str, ckpt_path: str):
     PREFIX_KEY = "net."
-    original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=weights_only))
-
-    if "Cosmos-1.0" in transformer_type:
-        TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0
-        TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0
-    elif "Cosmos-2.0" in transformer_type:
-        TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0
-        TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0
-    else:
-        assert False
+    original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True))
 
     with init_empty_weights():
         config = TRANSFORMER_CONFIGS[transformer_type]
@@ -398,69 +281,13 @@ def convert_vae(vae_type: str):
     return vae
 
 
-def save_pipeline_cosmos_1_0(args, transformer, vae):
-    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16)
-    tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path)
-    # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly.
-    # So, the sigma_min values that is used is the default value of 0.002.
-    scheduler = EDMEulerScheduler(
-        sigma_min=0.002,
-        sigma_max=80,
-        sigma_data=0.5,
-        sigma_schedule="karras",
-        num_train_timesteps=1000,
-        prediction_type="epsilon",
-        rho=7.0,
-        final_sigmas_type="sigma_min",
-    )
-
-    pipe_cls = CosmosTextToWorldPipeline if "Text2World" in args.transformer_type else CosmosVideoToWorldPipeline
-    pipe = pipe_cls(
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        transformer=transformer,
-        vae=vae,
-        scheduler=scheduler,
-    )
-    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
-
-
-def save_pipeline_cosmos_2_0(args, transformer, vae):
-    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16)
-    tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path)
-
-    scheduler = EDMEulerScheduler(
-        sigma_min=0.002,
-        sigma_max=80,
-        sigma_data=1.0,
-        sigma_schedule="karras",
-        num_train_timesteps=1000,
-        prediction_type="epsilon",
-        rho=7.0,
-        final_sigmas_type="sigma_min",
-        use_flow_sigmas=True,
-    )
-
-    pipe_cls = Cosmos2TextToImagePipeline if "Text2Image" in args.transformer_type else Cosmos2VideoToWorldPipeline
-    pipe = pipe_cls(
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        transformer=transformer,
-        vae=vae,
-        scheduler=scheduler,
-    )
-    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
-
-
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys()))
     parser.add_argument(
         "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
     )
-    parser.add_argument(
-        "--vae_type", type=str, default=None, choices=["none", *list(VAE_CONFIGS.keys())], help="Type of VAE"
-    )
+    parser.add_argument("--vae_type", type=str, default=None, choices=list(VAE_CONFIGS.keys()), help="Type of VAE")
     parser.add_argument("--text_encoder_path", type=str, default="google-t5/t5-11b")
     parser.add_argument("--tokenizer_path", type=str, default="google-t5/t5-11b")
     parser.add_argument("--save_pipeline", action="store_true")
@@ -489,26 +316,37 @@ if __name__ == "__main__":
         assert args.tokenizer_path is not None
 
     if args.transformer_ckpt_path is not None:
-        weights_only = "Cosmos-1.0" in args.transformer_type
-        transformer = convert_transformer(args.transformer_type, args.transformer_ckpt_path, weights_only)
+        transformer = convert_transformer(args.transformer_type, args.transformer_ckpt_path)
         transformer = transformer.to(dtype=dtype)
         if not args.save_pipeline:
             transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
 
     if args.vae_type is not None:
-        if "Cosmos-1.0" in args.transformer_type:
-            vae = convert_vae(args.vae_type)
-        else:
-            vae = AutoencoderKLWan.from_pretrained(
-                "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32
-            )
+        vae = convert_vae(args.vae_type)
         if not args.save_pipeline:
             vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")
 
     if args.save_pipeline:
-        if "Cosmos-1.0" in args.transformer_type:
-            save_pipeline_cosmos_1_0(args, transformer, vae)
-        elif "Cosmos-2.0" in args.transformer_type:
-            save_pipeline_cosmos_2_0(args, transformer, vae)
-        else:
-            assert False
+        text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=dtype)
+        tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path)
+        # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly.
+        # So, the sigma_min values that is used is the default value of 0.002.
+        scheduler = EDMEulerScheduler(
+            sigma_min=0.002,
+            sigma_max=80,
+            sigma_data=0.5,
+            sigma_schedule="karras",
+            num_train_timesteps=1000,
+            prediction_type="epsilon",
+            rho=7.0,
+            final_sigmas_type="sigma_min",
+        )
+
+        pipe = CosmosTextToWorldPipeline(
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            vae=vae,
+            scheduler=scheduler,
+        )
+        pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB")

scripts/convert_wan_to_diffusers.py

@@ -1,6 +1,6 @@
 import argparse
 import pathlib
-from typing import Any, Dict, Tuple
+from typing import Any, Dict
 
 import torch
 from accelerate import init_empty_weights
@@ -14,8 +14,6 @@ from diffusers import (
     WanImageToVideoPipeline,
     WanPipeline,
     WanTransformer3DModel,
-    WanVACEPipeline,
-    WanVACETransformer3DModel,
 )
 
 
@@ -61,52 +59,7 @@ TRANSFORMER_KEYS_RENAME_DICT = {
     "attn2.norm_k_img": "attn2.norm_added_k",
 }
 
-VACE_TRANSFORMER_KEYS_RENAME_DICT = {
-    "time_embedding.0": "condition_embedder.time_embedder.linear_1",
-    "time_embedding.2": "condition_embedder.time_embedder.linear_2",
-    "text_embedding.0": "condition_embedder.text_embedder.linear_1",
-    "text_embedding.2": "condition_embedder.text_embedder.linear_2",
-    "time_projection.1": "condition_embedder.time_proj",
-    "head.modulation": "scale_shift_table",
-    "head.head": "proj_out",
-    "modulation": "scale_shift_table",
-    "ffn.0": "ffn.net.0.proj",
-    "ffn.2": "ffn.net.2",
-    # Hack to swap the layer names
-    # The original model calls the norms in following order: norm1, norm3, norm2
-    # We convert it to: norm1, norm2, norm3
-    "norm2": "norm__placeholder",
-    "norm3": "norm2",
-    "norm__placeholder": "norm3",
-    # # For the I2V model
-    # "img_emb.proj.0": "condition_embedder.image_embedder.norm1",
-    # "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj",
-    # "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2",
-    # "img_emb.proj.4": "condition_embedder.image_embedder.norm2",
-    # # for the FLF2V model
-    # "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed",
-    # Add attention component mappings
-    "self_attn.q": "attn1.to_q",
-    "self_attn.k": "attn1.to_k",
-    "self_attn.v": "attn1.to_v",
-    "self_attn.o": "attn1.to_out.0",
-    "self_attn.norm_q": "attn1.norm_q",
-    "self_attn.norm_k": "attn1.norm_k",
-    "cross_attn.q": "attn2.to_q",
-    "cross_attn.k": "attn2.to_k",
-    "cross_attn.v": "attn2.to_v",
-    "cross_attn.o": "attn2.to_out.0",
-    "cross_attn.norm_q": "attn2.norm_q",
-    "cross_attn.norm_k": "attn2.norm_k",
-    "attn2.to_k_img": "attn2.add_k_proj",
-    "attn2.to_v_img": "attn2.add_v_proj",
-    "attn2.norm_k_img": "attn2.norm_added_k",
-    "before_proj": "proj_in",
-    "after_proj": "proj_out",
-}
-
 TRANSFORMER_SPECIAL_KEYS_REMAP = {}
-VACE_TRANSFORMER_SPECIAL_KEYS_REMAP = {}
 
 
 def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> Dict[str, Any]:
@@ -121,7 +74,7 @@ def load_sharded_safetensors(dir: pathlib.Path):
     return state_dict
 
 
-def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
+def get_transformer_config(model_type: str) -> Dict[str, Any]:
     if model_type == "Wan-T2V-1.3B":
         config = {
             "model_id": "StevenZhang/Wan2.1-T2V-1.3B-Diff",
@@ -141,8 +94,6 @@ def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
                 "text_dim": 4096,
             },
         }
-        RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP
     elif model_type == "Wan-T2V-14B":
         config = {
             "model_id": "StevenZhang/Wan2.1-T2V-14B-Diff",
@@ -162,8 +113,6 @@ def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
                 "text_dim": 4096,
             },
         }
-        RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP
     elif model_type == "Wan-I2V-14B-480p":
         config = {
             "model_id": "StevenZhang/Wan2.1-I2V-14B-480P-Diff",
@@ -184,8 +133,6 @@ def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
                 "text_dim": 4096,
             },
         }
-        RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP
     elif model_type == "Wan-I2V-14B-720p":
         config = {
             "model_id": "StevenZhang/Wan2.1-I2V-14B-720P-Diff",
@@ -206,8 +153,6 @@ def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
                 "text_dim": 4096,
             },
         }
-        RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP
     elif model_type == "Wan-FLF2V-14B-720P":
         config = {
             "model_id": "ypyp/Wan2.1-FLF2V-14B-720P",  # This is just a placeholder
@@ -230,60 +175,11 @@ def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]:
                 "pos_embed_seq_len": 257 * 2,
             },
         }
-        RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP
-    elif model_type == "Wan-VACE-1.3B":
-        config = {
-            "model_id": "Wan-AI/Wan2.1-VACE-1.3B",
-            "diffusers_config": {
-                "added_kv_proj_dim": None,
-                "attention_head_dim": 128,
-                "cross_attn_norm": True,
-                "eps": 1e-06,
-                "ffn_dim": 8960,
-                "freq_dim": 256,
-                "in_channels": 16,
-                "num_attention_heads": 12,
-                "num_layers": 30,
-                "out_channels": 16,
-                "patch_size": [1, 2, 2],
-                "qk_norm": "rms_norm_across_heads",
-                "text_dim": 4096,
-                "vace_layers": [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28],
-                "vace_in_channels": 96,
-            },
-        }
-        RENAME_DICT = VACE_TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = VACE_TRANSFORMER_SPECIAL_KEYS_REMAP
-    elif model_type == "Wan-VACE-14B":
-        config = {
-            "model_id": "Wan-AI/Wan2.1-VACE-14B",
-            "diffusers_config": {
-                "added_kv_proj_dim": None,
-                "attention_head_dim": 128,
-                "cross_attn_norm": True,
-                "eps": 1e-06,
-                "ffn_dim": 13824,
-                "freq_dim": 256,
-                "in_channels": 16,
-                "num_attention_heads": 40,
-                "num_layers": 40,
-                "out_channels": 16,
-                "patch_size": [1, 2, 2],
-                "qk_norm": "rms_norm_across_heads",
-                "text_dim": 4096,
-                "vace_layers": [0, 5, 10, 15, 20, 25, 30, 35],
-                "vace_in_channels": 96,
-            },
-        }
-        RENAME_DICT = VACE_TRANSFORMER_KEYS_RENAME_DICT
-        SPECIAL_KEYS_REMAP = VACE_TRANSFORMER_SPECIAL_KEYS_REMAP
-    return config, RENAME_DICT, SPECIAL_KEYS_REMAP
+    return config
 
 
 def convert_transformer(model_type: str):
-    config, RENAME_DICT, SPECIAL_KEYS_REMAP = get_transformer_config(model_type)
-
+    config = get_transformer_config(model_type)
     diffusers_config = config["diffusers_config"]
     model_id = config["model_id"]
     model_dir = pathlib.Path(snapshot_download(model_id, repo_type="model"))
@@ -291,19 +187,16 @@ def convert_transformer(model_type: str):
     original_state_dict = load_sharded_safetensors(model_dir)
 
     with init_empty_weights():
-        if "VACE" not in model_type:
-            transformer = WanTransformer3DModel.from_config(diffusers_config)
-        else:
-            transformer = WanVACETransformer3DModel.from_config(diffusers_config)
+        transformer = WanTransformer3DModel.from_config(diffusers_config)
 
     for key in list(original_state_dict.keys()):
         new_key = key[:]
-        for replace_key, rename_key in RENAME_DICT.items():
+        for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items():
             new_key = new_key.replace(replace_key, rename_key)
         update_state_dict_(original_state_dict, key, new_key)
 
     for key in list(original_state_dict.keys()):
-        for special_key, handler_fn_inplace in SPECIAL_KEYS_REMAP.items():
+        for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items():
             if special_key not in key:
                 continue
             handler_fn_inplace(key, original_state_dict)
@@ -519,7 +412,7 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_type", type=str, default=None)
     parser.add_argument("--output_path", type=str, required=True)
-    parser.add_argument("--dtype", default="fp32", choices=["fp32", "fp16", "bf16", "none"])
+    parser.add_argument("--dtype", default="fp32")
     return parser.parse_args()
 
 
@@ -533,20 +426,18 @@ DTYPE_MAPPING = {
 if __name__ == "__main__":
     args = get_args()
 
-    transformer = convert_transformer(args.model_type)
+    transformer = None
+    dtype = DTYPE_MAPPING[args.dtype]
+
+    transformer = convert_transformer(args.model_type).to(dtype=dtype)
     vae = convert_vae()
-    text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16)
+    text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl")
     tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl")
     flow_shift = 16.0 if "FLF2V" in args.model_type else 3.0
     scheduler = UniPCMultistepScheduler(
         prediction_type="flow_prediction", use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=flow_shift
     )
 
-    # If user has specified "none", we keep the original dtypes of the state dict without any conversion
-    if args.dtype != "none":
-        dtype = DTYPE_MAPPING[args.dtype]
-        transformer.to(dtype)
-
     if "I2V" in args.model_type or "FLF2V" in args.model_type:
         image_encoder = CLIPVisionModelWithProjection.from_pretrained(
             "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=torch.bfloat16
@@ -561,14 +452,6 @@ if __name__ == "__main__":
             image_encoder=image_encoder,
             image_processor=image_processor,
         )
-    elif "VACE" in args.model_type:
-        pipe = WanVACEPipeline(
-            transformer=transformer,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            vae=vae,
-            scheduler=scheduler,
-        )
     else:
         pipe = WanPipeline(
             transformer=transformer,

src/diffusers/__init__.py

@@ -215,7 +215,6 @@ else:
             "UVit2DModel",
             "VQModel",
             "WanTransformer3DModel",
-            "WanVACETransformer3DModel",
         ]
     )
     _import_structure["optimization"] = [
@@ -361,8 +360,6 @@ else:
             "CogView4ControlPipeline",
             "CogView4Pipeline",
             "ConsisIDPipeline",
-            "Cosmos2TextToImagePipeline",
-            "Cosmos2VideoToWorldPipeline",
             "CosmosTextToWorldPipeline",
             "CosmosVideoToWorldPipeline",
             "CycleDiffusionPipeline",
@@ -444,7 +441,6 @@ else:
             "SanaControlNetPipeline",
             "SanaPAGPipeline",
             "SanaPipeline",
-            "SanaSprintImg2ImgPipeline",
             "SanaSprintPipeline",
             "SemanticStableDiffusionPipeline",
             "ShapEImg2ImgPipeline",
@@ -530,7 +526,6 @@ else:
             "VQDiffusionPipeline",
             "WanImageToVideoPipeline",
             "WanPipeline",
-            "WanVACEPipeline",
             "WanVideoToVideoPipeline",
             "WuerstchenCombinedPipeline",
             "WuerstchenDecoderPipeline",
@@ -696,7 +691,6 @@ else:
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     from .configuration_utils import ConfigMixin
-    from .quantizers import PipelineQuantizationConfig
 
     try:
         if not is_bitsandbytes_available():
@@ -825,7 +819,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             UVit2DModel,
             VQModel,
             WanTransformer3DModel,
-            WanVACETransformer3DModel,
         )
         from .optimization import (
             get_constant_schedule,
@@ -951,8 +944,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             CogView4ControlPipeline,
             CogView4Pipeline,
             ConsisIDPipeline,
-            Cosmos2TextToImagePipeline,
-            Cosmos2VideoToWorldPipeline,
             CosmosTextToWorldPipeline,
             CosmosVideoToWorldPipeline,
             CycleDiffusionPipeline,
@@ -1034,7 +1025,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             SanaControlNetPipeline,
             SanaPAGPipeline,
             SanaPipeline,
-            SanaSprintImg2ImgPipeline,
             SanaSprintPipeline,
             SemanticStableDiffusionPipeline,
             ShapEImg2ImgPipeline,
@@ -1119,7 +1109,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             VQDiffusionPipeline,
             WanImageToVideoPipeline,
             WanPipeline,
-            WanVACEPipeline,
             WanVideoToVideoPipeline,
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,

src/diffusers/hooks/faster_cache.py

@@ -146,7 +146,7 @@ class FasterCacheConfig:
     alpha_low_frequency: float = 1.1
     alpha_high_frequency: float = 1.1
 
-    # n as described in CFG-Cache explanation in the paper - dependent on the model
+    # n as described in CFG-Cache explanation in the paper - dependant on the model
     unconditional_batch_skip_range: int = 5
     unconditional_batch_timestep_skip_range: Tuple[int, int] = (-1, 641)
 

src/diffusers/hooks/group_offloading.py

@@ -113,7 +113,6 @@ class ModuleGroup:
         finally:
             pinned_dict = None
 
-    @torch.compiler.disable()
     def onload_(self):
         r"""Onloads the group of modules to the onload_device."""
         torch_accelerator_module = (
@@ -166,7 +165,6 @@ class ModuleGroup:
                     if self.record_stream:
                         buffer.data.record_stream(current_stream)
 
-    @torch.compiler.disable()
     def offload_(self):
         r"""Offloads the group of modules to the offload_device."""
 

src/diffusers/hooks/hooks.py

@@ -45,7 +45,7 @@ class ModelHook:
 
     def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         r"""
-        Hook that is executed when a model is deinitialized.
+        Hook that is executed when a model is deinitalized.
 
         Args:
             module (`torch.nn.Module`):

src/diffusers/hooks/layerwise_casting.py

@@ -62,7 +62,7 @@ class LayerwiseCastingHook(ModelHook):
 
     def deinitalize_hook(self, module: torch.nn.Module):
         raise NotImplementedError(
-            "LayerwiseCastingHook does not support deinitialization. A model once enabled with layerwise casting will "
+            "LayerwiseCastingHook does not support deinitalization. A model once enabled with layerwise casting will "
             "have casted its weights to a lower precision dtype for storage. Casting this back to the original dtype "
             "will lead to precision loss, which might have an impact on the model's generation quality. The model should "
             "be re-initialized and loaded in the original dtype."

src/diffusers/loaders/lora_base.py

@@ -465,7 +465,7 @@ class LoraBaseMixin:
     """Utility class for handling LoRAs."""
 
     _lora_loadable_modules = []
-    _merged_adapters = set()
+    num_fused_loras = 0
 
     def load_lora_weights(self, **kwargs):
         raise NotImplementedError("`load_lora_weights()` is not implemented.")
@@ -592,9 +592,6 @@ class LoraBaseMixin:
         if len(components) == 0:
             raise ValueError("`components` cannot be an empty list.")
 
-        # Need to retrieve the names as `adapter_names` can be None. So we cannot directly use it
-        # in `self._merged_adapters = self._merged_adapters | merged_adapter_names`.
-        merged_adapter_names = set()
         for fuse_component in components:
             if fuse_component not in self._lora_loadable_modules:
                 raise ValueError(f"{fuse_component} is not found in {self._lora_loadable_modules=}.")
@@ -604,19 +601,13 @@ class LoraBaseMixin:
                 # check if diffusers model
                 if issubclass(model.__class__, ModelMixin):
                     model.fuse_lora(lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names)
-                    for module in model.modules():
-                        if isinstance(module, BaseTunerLayer):
-                            merged_adapter_names.update(set(module.merged_adapters))
                 # handle transformers models.
                 if issubclass(model.__class__, PreTrainedModel):
                     fuse_text_encoder_lora(
                         model, lora_scale=lora_scale, safe_fusing=safe_fusing, adapter_names=adapter_names
                     )
-                    for module in model.modules():
-                        if isinstance(module, BaseTunerLayer):
-                            merged_adapter_names.update(set(module.merged_adapters))
 
-        self._merged_adapters = self._merged_adapters | merged_adapter_names
+        self.num_fused_loras += 1
 
     def unfuse_lora(self, components: List[str] = [], **kwargs):
         r"""
@@ -670,18 +661,9 @@ class LoraBaseMixin:
                 if issubclass(model.__class__, (ModelMixin, PreTrainedModel)):
                     for module in model.modules():
                         if isinstance(module, BaseTunerLayer):
-                            for adapter in set(module.merged_adapters):
-                                if adapter and adapter in self._merged_adapters:
-                                    self._merged_adapters = self._merged_adapters - {adapter}
                             module.unmerge()
 
-    @property
-    def num_fused_loras(self):
-        return len(self._merged_adapters)
-
-    @property
-    def fused_loras(self):
-        return self._merged_adapters
+        self.num_fused_loras -= 1
 
     def set_adapters(
         self,

src/diffusers/loaders/peft.py

@@ -58,7 +58,6 @@ _SET_ADAPTER_SCALE_FN_MAPPING = {
     "CogView4Transformer2DModel": lambda model_cls, weights: weights,
     "HiDreamImageTransformer2DModel": lambda model_cls, weights: weights,
     "HunyuanVideoFramepackTransformer3DModel": lambda model_cls, weights: weights,
-    "WanVACETransformer3DModel": lambda model_cls, weights: weights,
 }
 
 
@@ -252,7 +251,7 @@ class PeftAdapterMixin:
 
             rank = {}
             for key, val in state_dict.items():
-                # Cannot figure out rank from lora layers that don't have at least 2 dimensions.
+                # Cannot figure out rank from lora layers that don't have atleast 2 dimensions.
                 # Bias layers in LoRA only have a single dimension
                 if "lora_B" in key and val.ndim > 1:
                     # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol.

src/diffusers/models/__init__.py

@@ -89,7 +89,6 @@ if is_torch_available():
     _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"]
     _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"]
     _import_structure["transformers.transformer_wan"] = ["WanTransformer3DModel"]
-    _import_structure["transformers.transformer_wan_vace"] = ["WanVACETransformer3DModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
@@ -179,7 +178,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             Transformer2DModel,
             TransformerTemporalModel,
             WanTransformer3DModel,
-            WanVACETransformer3DModel,
         )
         from .unets import (
             I2VGenXLUNet,

src/diffusers/models/autoencoders/autoencoder_kl.py

@@ -63,8 +63,8 @@ class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin, PeftAdapter
             Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
         force_upcast (`bool`, *optional*, default to `True`):
             If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
-            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
-            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
         mid_block_add_attention (`bool`, *optional*, default to `True`):
             If enabled, the mid_block of the Encoder and Decoder will have attention blocks. If set to false, the
             mid_block will only have resnet blocks

src/diffusers/models/autoencoders/autoencoder_kl_allegro.py

@@ -715,8 +715,8 @@ class AutoencoderKLAllegro(ModelMixin, ConfigMixin):
             Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
         force_upcast (`bool`, default to `True`):
             If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
-            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
-            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
     """
 
     _supports_gradient_checkpointing = True

src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py

@@ -983,8 +983,8 @@ class AutoencoderKLCogVideoX(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
         force_upcast (`bool`, *optional*, default to `True`):
             If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
-            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
-            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
     """
 
     _supports_gradient_checkpointing = True

src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py

@@ -161,8 +161,8 @@ class AutoencoderKLTemporalDecoder(ModelMixin, ConfigMixin):
             Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752) paper.
         force_upcast (`bool`, *optional*, default to `True`):
             If enabled it will force the VAE to run in float32 for high image resolution pipelines, such as SD-XL. VAE
-            can be fine-tuned / trained to a lower range without losing too much precision in which case `force_upcast`
-            can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
+            can be fine-tuned / trained to a lower range without loosing too much precision in which case
+            `force_upcast` can be set to `False` - see: https://huggingface.co/madebyollin/sdxl-vae-fp16-fix
     """
 
     _supports_gradient_checkpointing = True

src/diffusers/models/transformers/__init__.py

@@ -32,4 +32,3 @@ if is_torch_available():
     from .transformer_sd3 import SD3Transformer2DModel
     from .transformer_temporal import TransformerTemporalModel
     from .transformer_wan import WanTransformer3DModel
-    from .transformer_wan_vace import WanVACETransformer3DModel

src/diffusers/models/transformers/transformer_cosmos.py

@@ -100,15 +100,11 @@ class CosmosAdaLayerNorm(nn.Module):
         embedded_timestep = self.linear_2(embedded_timestep)
 
         if temb is not None:
-            embedded_timestep = embedded_timestep + temb[..., : 2 * self.embedding_dim]
+            embedded_timestep = embedded_timestep + temb[:, : 2 * self.embedding_dim]
 
-        shift, scale = embedded_timestep.chunk(2, dim=-1)
+        shift, scale = embedded_timestep.chunk(2, dim=1)
         hidden_states = self.norm(hidden_states)
-
-        if embedded_timestep.ndim == 2:
-            shift, scale = (x.unsqueeze(1) for x in (shift, scale))
-
-        hidden_states = hidden_states * (1 + scale) + shift
+        hidden_states = hidden_states * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         return hidden_states
 
 
@@ -139,13 +135,9 @@ class CosmosAdaLayerNormZero(nn.Module):
         if temb is not None:
             embedded_timestep = embedded_timestep + temb
 
-        shift, scale, gate = embedded_timestep.chunk(3, dim=-1)
+        shift, scale, gate = embedded_timestep.chunk(3, dim=1)
         hidden_states = self.norm(hidden_states)
-
-        if embedded_timestep.ndim == 2:
-            shift, scale, gate = (x.unsqueeze(1) for x in (shift, scale, gate))
-
-        hidden_states = hidden_states * (1 + scale) + shift
+        hidden_states = hidden_states * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
         return hidden_states, gate
 
 
@@ -263,19 +255,19 @@ class CosmosTransformerBlock(nn.Module):
         # 1. Self Attention
         norm_hidden_states, gate = self.norm1(hidden_states, embedded_timestep, temb)
         attn_output = self.attn1(norm_hidden_states, image_rotary_emb=image_rotary_emb)
-        hidden_states = hidden_states + gate * attn_output
+        hidden_states = hidden_states + gate.unsqueeze(1) * attn_output
 
         # 2. Cross Attention
         norm_hidden_states, gate = self.norm2(hidden_states, embedded_timestep, temb)
         attn_output = self.attn2(
             norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask
         )
-        hidden_states = hidden_states + gate * attn_output
+        hidden_states = hidden_states + gate.unsqueeze(1) * attn_output
 
         # 3. Feed Forward
         norm_hidden_states, gate = self.norm3(hidden_states, embedded_timestep, temb)
         ff_output = self.ff(norm_hidden_states)
-        hidden_states = hidden_states + gate * ff_output
+        hidden_states = hidden_states + gate.unsqueeze(1) * ff_output
 
         return hidden_states
 
@@ -521,23 +513,7 @@ class CosmosTransformer3DModel(ModelMixin, ConfigMixin):
         hidden_states = hidden_states.flatten(1, 3)  # [B, T, H, W, C] -> [B, THW, C]
 
         # 4. Timestep embeddings
-        if timestep.ndim == 1:
-            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
-        elif timestep.ndim == 5:
-            assert timestep.shape == (batch_size, 1, num_frames, 1, 1), (
-                f"Expected timestep to have shape [B, 1, T, 1, 1], but got {timestep.shape}"
-            )
-            timestep = timestep.flatten()
-            temb, embedded_timestep = self.time_embed(hidden_states, timestep)
-            # We can do this because num_frames == post_patch_num_frames, as p_t is 1
-            temb, embedded_timestep = (
-                x.view(batch_size, post_patch_num_frames, 1, 1, -1)
-                .expand(-1, -1, post_patch_height, post_patch_width, -1)
-                .flatten(1, 3)
-                for x in (temb, embedded_timestep)
-            )  # [BT, C] -> [B, T, 1, 1, C] -> [B, T, H, W, C] -> [B, THW, C]
-        else:
-            assert False
+        temb, embedded_timestep = self.time_embed(hidden_states, timestep)
 
         # 5. Transformer blocks
         for block in self.transformer_blocks:
@@ -568,6 +544,8 @@ class CosmosTransformer3DModel(ModelMixin, ConfigMixin):
         hidden_states = self.proj_out(hidden_states)
         hidden_states = hidden_states.unflatten(2, (p_h, p_w, p_t, -1))
         hidden_states = hidden_states.unflatten(1, (post_patch_num_frames, post_patch_height, post_patch_width))
+        # Please just kill me at this point. What even is this permutation order and why is it different from the patching order?
+        # Another few hours of sanity lost to the void.
         hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
         hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
 

src/diffusers/models/transformers/transformer_flux.py

@@ -447,6 +447,8 @@ class FluxTransformer2DModel(
         timestep = timestep.to(hidden_states.dtype) * 1000
         if guidance is not None:
             guidance = guidance.to(hidden_states.dtype) * 1000
+        else:
+            guidance = None
 
         temb = (
             self.time_text_embed(timestep, pooled_projections)

src/diffusers/models/transformers/transformer_wan.py

@@ -72,8 +72,7 @@ class WanAttnProcessor2_0:
         if rotary_emb is not None:
 
             def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor):
-                dtype = torch.float32 if hidden_states.device.type == "mps" else torch.float64
-                x_rotated = torch.view_as_complex(hidden_states.to(dtype).unflatten(3, (-1, 2)))
+                x_rotated = torch.view_as_complex(hidden_states.to(torch.float64).unflatten(3, (-1, 2)))
                 x_out = torch.view_as_real(x_rotated * freqs).flatten(3, 4)
                 return x_out.type_as(hidden_states)
 
@@ -191,10 +190,9 @@ class WanRotaryPosEmbed(nn.Module):
         t_dim = attention_head_dim - h_dim - w_dim
 
         freqs = []
-        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
         for dim in [t_dim, h_dim, w_dim]:
             freq = get_1d_rotary_pos_embed(
-                dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=freqs_dtype
+                dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float64
             )
             freqs.append(freq)
         self.freqs = torch.cat(freqs, dim=1)

src/diffusers/models/transformers/transformer_wan_vace.py

@@ -1,393 +0,0 @@
-# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
-#
-# 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 math
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-
-from ...configuration_utils import ConfigMixin, register_to_config
-from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
-from ...utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
-from ..attention import FeedForward
-from ..attention_processor import Attention
-from ..cache_utils import CacheMixin
-from ..modeling_outputs import Transformer2DModelOutput
-from ..modeling_utils import ModelMixin
-from ..normalization import FP32LayerNorm
-from .transformer_wan import WanAttnProcessor2_0, WanRotaryPosEmbed, WanTimeTextImageEmbedding, WanTransformerBlock
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class WanVACETransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        ffn_dim: int,
-        num_heads: int,
-        qk_norm: str = "rms_norm_across_heads",
-        cross_attn_norm: bool = False,
-        eps: float = 1e-6,
-        added_kv_proj_dim: Optional[int] = None,
-        apply_input_projection: bool = False,
-        apply_output_projection: bool = False,
-    ):
-        super().__init__()
-
-        # 1. Input projection
-        self.proj_in = None
-        if apply_input_projection:
-            self.proj_in = nn.Linear(dim, dim)
-
-        # 2. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.attn1 = Attention(
-            query_dim=dim,
-            heads=num_heads,
-            kv_heads=num_heads,
-            dim_head=dim // num_heads,
-            qk_norm=qk_norm,
-            eps=eps,
-            bias=True,
-            cross_attention_dim=None,
-            out_bias=True,
-            processor=WanAttnProcessor2_0(),
-        )
-
-        # 3. Cross-attention
-        self.attn2 = Attention(
-            query_dim=dim,
-            heads=num_heads,
-            kv_heads=num_heads,
-            dim_head=dim // num_heads,
-            qk_norm=qk_norm,
-            eps=eps,
-            bias=True,
-            cross_attention_dim=None,
-            out_bias=True,
-            added_kv_proj_dim=added_kv_proj_dim,
-            added_proj_bias=True,
-            processor=WanAttnProcessor2_0(),
-        )
-        self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
-
-        # 4. Feed-forward
-        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
-        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-
-        # 5. Output projection
-        self.proj_out = None
-        if apply_output_projection:
-            self.proj_out = nn.Linear(dim, dim)
-
-        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        control_hidden_states: torch.Tensor,
-        temb: torch.Tensor,
-        rotary_emb: torch.Tensor,
-    ) -> torch.Tensor:
-        if self.proj_in is not None:
-            control_hidden_states = self.proj_in(control_hidden_states)
-            control_hidden_states = control_hidden_states + hidden_states
-
-        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-            self.scale_shift_table + temb.float()
-        ).chunk(6, dim=1)
-
-        # 1. Self-attention
-        norm_hidden_states = (self.norm1(control_hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(
-            control_hidden_states
-        )
-        attn_output = self.attn1(hidden_states=norm_hidden_states, rotary_emb=rotary_emb)
-        control_hidden_states = (control_hidden_states.float() + attn_output * gate_msa).type_as(control_hidden_states)
-
-        # 2. Cross-attention
-        norm_hidden_states = self.norm2(control_hidden_states.float()).type_as(control_hidden_states)
-        attn_output = self.attn2(hidden_states=norm_hidden_states, encoder_hidden_states=encoder_hidden_states)
-        control_hidden_states = control_hidden_states + attn_output
-
-        # 3. Feed-forward
-        norm_hidden_states = (self.norm3(control_hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as(
-            control_hidden_states
-        )
-        ff_output = self.ffn(norm_hidden_states)
-        control_hidden_states = (control_hidden_states.float() + ff_output.float() * c_gate_msa).type_as(
-            control_hidden_states
-        )
-
-        conditioning_states = None
-        if self.proj_out is not None:
-            conditioning_states = self.proj_out(control_hidden_states)
-
-        return conditioning_states, control_hidden_states
-
-
-class WanVACETransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin):
-    r"""
-    A Transformer model for video-like data used in the Wan model.
-
-    Args:
-        patch_size (`Tuple[int]`, defaults to `(1, 2, 2)`):
-            3D patch dimensions for video embedding (t_patch, h_patch, w_patch).
-        num_attention_heads (`int`, defaults to `40`):
-            Fixed length for text embeddings.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of channels in each head.
-        in_channels (`int`, defaults to `16`):
-            The number of channels in the input.
-        out_channels (`int`, defaults to `16`):
-            The number of channels in the output.
-        text_dim (`int`, defaults to `512`):
-            Input dimension for text embeddings.
-        freq_dim (`int`, defaults to `256`):
-            Dimension for sinusoidal time embeddings.
-        ffn_dim (`int`, defaults to `13824`):
-            Intermediate dimension in feed-forward network.
-        num_layers (`int`, defaults to `40`):
-            The number of layers of transformer blocks to use.
-        window_size (`Tuple[int]`, defaults to `(-1, -1)`):
-            Window size for local attention (-1 indicates global attention).
-        cross_attn_norm (`bool`, defaults to `True`):
-            Enable cross-attention normalization.
-        qk_norm (`bool`, defaults to `True`):
-            Enable query/key normalization.
-        eps (`float`, defaults to `1e-6`):
-            Epsilon value for normalization layers.
-        add_img_emb (`bool`, defaults to `False`):
-            Whether to use img_emb.
-        added_kv_proj_dim (`int`, *optional*, defaults to `None`):
-            The number of channels to use for the added key and value projections. If `None`, no projection is used.
-    """
-
-    _supports_gradient_checkpointing = True
-    _skip_layerwise_casting_patterns = ["patch_embedding", "vace_patch_embedding", "condition_embedder", "norm"]
-    _no_split_modules = ["WanTransformerBlock", "WanVACETransformerBlock"]
-    _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"]
-    _keys_to_ignore_on_load_unexpected = ["norm_added_q"]
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: Tuple[int] = (1, 2, 2),
-        num_attention_heads: int = 40,
-        attention_head_dim: int = 128,
-        in_channels: int = 16,
-        out_channels: int = 16,
-        text_dim: int = 4096,
-        freq_dim: int = 256,
-        ffn_dim: int = 13824,
-        num_layers: int = 40,
-        cross_attn_norm: bool = True,
-        qk_norm: Optional[str] = "rms_norm_across_heads",
-        eps: float = 1e-6,
-        image_dim: Optional[int] = None,
-        added_kv_proj_dim: Optional[int] = None,
-        rope_max_seq_len: int = 1024,
-        pos_embed_seq_len: Optional[int] = None,
-        vace_layers: List[int] = [0, 5, 10, 15, 20, 25, 30, 35],
-        vace_in_channels: int = 96,
-    ) -> None:
-        super().__init__()
-
-        inner_dim = num_attention_heads * attention_head_dim
-        out_channels = out_channels or in_channels
-
-        if max(vace_layers) >= num_layers:
-            raise ValueError(f"VACE layers {vace_layers} exceed the number of transformer layers {num_layers}.")
-        if 0 not in vace_layers:
-            raise ValueError("VACE layers must include layer 0.")
-
-        # 1. Patch & position embedding
-        self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
-        self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size)
-        self.vace_patch_embedding = nn.Conv3d(vace_in_channels, inner_dim, kernel_size=patch_size, stride=patch_size)
-
-        # 2. Condition embeddings
-        # image_embedding_dim=1280 for I2V model
-        self.condition_embedder = WanTimeTextImageEmbedding(
-            dim=inner_dim,
-            time_freq_dim=freq_dim,
-            time_proj_dim=inner_dim * 6,
-            text_embed_dim=text_dim,
-            image_embed_dim=image_dim,
-            pos_embed_seq_len=pos_embed_seq_len,
-        )
-
-        # 3. Transformer blocks
-        self.blocks = nn.ModuleList(
-            [
-                WanTransformerBlock(
-                    inner_dim, ffn_dim, num_attention_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim
-                )
-                for _ in range(num_layers)
-            ]
-        )
-
-        self.vace_blocks = nn.ModuleList(
-            [
-                WanVACETransformerBlock(
-                    inner_dim,
-                    ffn_dim,
-                    num_attention_heads,
-                    qk_norm,
-                    cross_attn_norm,
-                    eps,
-                    added_kv_proj_dim,
-                    apply_input_projection=i == 0,  # Layer 0 always has input projection and is in vace_layers
-                    apply_output_projection=True,
-                )
-                for i in range(len(vace_layers))
-            ]
-        )
-
-        # 4. Output norm & projection
-        self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False)
-        self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
-        self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5)
-
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_hidden_states_image: Optional[torch.Tensor] = None,
-        control_hidden_states: torch.Tensor = None,
-        control_hidden_states_scale: torch.Tensor = None,
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
-                logger.warning(
-                    "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
-                )
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p_t, p_h, p_w = self.config.patch_size
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p_h
-        post_patch_width = width // p_w
-
-        if control_hidden_states_scale is None:
-            control_hidden_states_scale = control_hidden_states.new_ones(len(self.config.vace_layers))
-        control_hidden_states_scale = torch.unbind(control_hidden_states_scale)
-        if len(control_hidden_states_scale) != len(self.config.vace_layers):
-            raise ValueError(
-                f"Length of `control_hidden_states_scale` {len(control_hidden_states_scale)} should be "
-                f"equal to {len(self.config.vace_layers)}."
-            )
-
-        # 1. Rotary position embedding
-        rotary_emb = self.rope(hidden_states)
-
-        # 2. Patch embedding
-        hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
-
-        control_hidden_states = self.vace_patch_embedding(control_hidden_states)
-        control_hidden_states = control_hidden_states.flatten(2).transpose(1, 2)
-        control_hidden_states_padding = control_hidden_states.new_zeros(
-            batch_size, hidden_states.size(1) - control_hidden_states.size(1), control_hidden_states.size(2)
-        )
-        control_hidden_states = torch.cat([control_hidden_states, control_hidden_states_padding], dim=1)
-
-        # 3. Time embedding
-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image
-        )
-        timestep_proj = timestep_proj.unflatten(1, (6, -1))
-
-        # 4. Image embedding
-        if encoder_hidden_states_image is not None:
-            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
-
-        # 5. Transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            # Prepare VACE hints
-            control_hidden_states_list = []
-            for i, block in enumerate(self.vace_blocks):
-                conditioning_states, control_hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states, control_hidden_states, timestep_proj, rotary_emb
-                )
-                control_hidden_states_list.append((conditioning_states, control_hidden_states_scale[i]))
-            control_hidden_states_list = control_hidden_states_list[::-1]
-
-            for i, block in enumerate(self.blocks):
-                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states, timestep_proj, rotary_emb
-                )
-                if i in self.config.vace_layers:
-                    control_hint, scale = control_hidden_states_list.pop()
-                    hidden_states = hidden_states + control_hint * scale
-        else:
-            # Prepare VACE hints
-            control_hidden_states_list = []
-            for i, block in enumerate(self.vace_blocks):
-                conditioning_states, control_hidden_states = block(
-                    hidden_states, encoder_hidden_states, control_hidden_states, timestep_proj, rotary_emb
-                )
-                control_hidden_states_list.append((conditioning_states, control_hidden_states_scale[i]))
-            control_hidden_states_list = control_hidden_states_list[::-1]
-
-            for i, block in enumerate(self.blocks):
-                hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
-                if i in self.config.vace_layers:
-                    control_hint, scale = control_hidden_states_list.pop()
-                    hidden_states = hidden_states + control_hint * scale
-
-        # 6. Output norm, projection & unpatchify
-        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
-
-        # Move the shift and scale tensors to the same device as hidden_states.
-        # When using multi-GPU inference via accelerate these will be on the
-        # first device rather than the last device, which hidden_states ends up
-        # on.
-        shift = shift.to(hidden_states.device)
-        scale = scale.to(hidden_states.device)
-
-        hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states)
-        hidden_states = self.proj_out(hidden_states)
-
-        hidden_states = hidden_states.reshape(
-            batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1
-        )
-        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
-        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not return_dict:
-            return (output,)
-
-        return Transformer2DModelOutput(sample=output)

src/diffusers/pipelines/__init__.py

@@ -157,12 +157,7 @@ else:
     _import_structure["cogview3"] = ["CogView3PlusPipeline"]
     _import_structure["cogview4"] = ["CogView4Pipeline", "CogView4ControlPipeline"]
     _import_structure["consisid"] = ["ConsisIDPipeline"]
-    _import_structure["cosmos"] = [
-        "Cosmos2TextToImagePipeline",
-        "CosmosTextToWorldPipeline",
-        "CosmosVideoToWorldPipeline",
-        "Cosmos2VideoToWorldPipeline",
-    ]
+    _import_structure["cosmos"] = ["CosmosTextToWorldPipeline", "CosmosVideoToWorldPipeline"]
     _import_structure["controlnet"].extend(
         [
             "BlipDiffusionControlNetPipeline",
@@ -295,12 +290,7 @@ else:
     _import_structure["paint_by_example"] = ["PaintByExamplePipeline"]
     _import_structure["pia"] = ["PIAPipeline"]
     _import_structure["pixart_alpha"] = ["PixArtAlphaPipeline", "PixArtSigmaPipeline"]
-    _import_structure["sana"] = [
-        "SanaPipeline",
-        "SanaSprintPipeline",
-        "SanaControlNetPipeline",
-        "SanaSprintImg2ImgPipeline",
-    ]
+    _import_structure["sana"] = ["SanaPipeline", "SanaSprintPipeline", "SanaControlNetPipeline"]
     _import_structure["semantic_stable_diffusion"] = ["SemanticStableDiffusionPipeline"]
     _import_structure["shap_e"] = ["ShapEImg2ImgPipeline", "ShapEPipeline"]
     _import_structure["stable_audio"] = [
@@ -376,7 +366,7 @@ else:
         "WuerstchenDecoderPipeline",
         "WuerstchenPriorPipeline",
     ]
-    _import_structure["wan"] = ["WanPipeline", "WanImageToVideoPipeline", "WanVideoToVideoPipeline", "WanVACEPipeline"]
+    _import_structure["wan"] = ["WanPipeline", "WanImageToVideoPipeline", "WanVideoToVideoPipeline"]
 try:
     if not is_onnx_available():
         raise OptionalDependencyNotAvailable()
@@ -564,12 +554,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             StableDiffusionControlNetXSPipeline,
             StableDiffusionXLControlNetXSPipeline,
         )
-        from .cosmos import (
-            Cosmos2TextToImagePipeline,
-            Cosmos2VideoToWorldPipeline,
-            CosmosTextToWorldPipeline,
-            CosmosVideoToWorldPipeline,
-        )
+        from .cosmos import CosmosTextToWorldPipeline, CosmosVideoToWorldPipeline
         from .deepfloyd_if import (
             IFImg2ImgPipeline,
             IFImg2ImgSuperResolutionPipeline,
@@ -690,7 +675,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
         from .paint_by_example import PaintByExamplePipeline
         from .pia import PIAPipeline
         from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline
-        from .sana import SanaControlNetPipeline, SanaPipeline, SanaSprintImg2ImgPipeline, SanaSprintPipeline
+        from .sana import SanaControlNetPipeline, SanaPipeline, SanaSprintPipeline
         from .semantic_stable_diffusion import SemanticStableDiffusionPipeline
         from .shap_e import ShapEImg2ImgPipeline, ShapEPipeline
         from .stable_audio import StableAudioPipeline, StableAudioProjectionModel
@@ -749,7 +734,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             UniDiffuserTextDecoder,
         )
         from .visualcloze import VisualClozeGenerationPipeline, VisualClozePipeline
-        from .wan import WanImageToVideoPipeline, WanPipeline, WanVACEPipeline, WanVideoToVideoPipeline
+        from .wan import WanImageToVideoPipeline, WanPipeline, WanVideoToVideoPipeline
         from .wuerstchen import (
             WuerstchenCombinedPipeline,
             WuerstchenDecoderPipeline,

src/diffusers/pipelines/consisid/consisid_utils.py

@@ -166,7 +166,7 @@ def process_face_embeddings(
         raise RuntimeError("facexlib align face fail")
     align_face = face_helper_1.cropped_faces[0]  # (512, 512, 3)  # RGB
 
-    # in case insightface didn't detect face
+    # incase insightface didn't detect face
     if id_ante_embedding is None:
         logger.warning("Failed to detect face using insightface. Extracting embedding with align face")
         id_ante_embedding = face_helper_2.get_feat(align_face)

src/diffusers/pipelines/cosmos/__init__.py

@@ -22,8 +22,6 @@ except OptionalDependencyNotAvailable:
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
-    _import_structure["pipeline_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"]
-    _import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"]
     _import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"]
     _import_structure["pipeline_cosmos_video2world"] = ["CosmosVideoToWorldPipeline"]
 
@@ -35,8 +33,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *
     else:
-        from .pipeline_cosmos2_text2image import Cosmos2TextToImagePipeline
-        from .pipeline_cosmos2_video2world import Cosmos2VideoToWorldPipeline
         from .pipeline_cosmos_text2world import CosmosTextToWorldPipeline
         from .pipeline_cosmos_video2world import CosmosVideoToWorldPipeline
 

src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py

@@ -1,653 +0,0 @@
-# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
-#
-# 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 inspect
-from typing import Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from transformers import T5EncoderModel, T5TokenizerFast
-
-from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...models import AutoencoderKLWan, CosmosTransformer3DModel
-from ...schedulers import EDMEulerScheduler
-from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring
-from ...utils.torch_utils import randn_tensor
-from ...video_processor import VideoProcessor
-from ..pipeline_utils import DiffusionPipeline
-from .pipeline_output import CosmosImagePipelineOutput
-
-
-if is_cosmos_guardrail_available():
-    from cosmos_guardrail import CosmosSafetyChecker
-else:
-
-    class CosmosSafetyChecker:
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                "`cosmos_guardrail` is not installed. Please install it to use the safety checker for Cosmos: `pip install cosmos_guardrail`."
-            )
-
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```python
-        >>> import torch
-        >>> from diffusers import Cosmos2TextToImagePipeline
-
-        >>> # Available checkpoints: nvidia/Cosmos-Predict2-2B-Text2Image, nvidia/Cosmos-Predict2-14B-Text2Image
-        >>> model_id = "nvidia/Cosmos-Predict2-2B-Text2Image"
-        >>> pipe = Cosmos2TextToImagePipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-
-        >>> prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
-        >>> negative_prompt = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
-
-        >>> output = pipe(
-        ...     prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator().manual_seed(1)
-        ... ).images[0]
-        >>> output.save("output.png")
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class Cosmos2TextToImagePipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-image generation using [Cosmos](https://github.com/NVIDIA/Cosmos).
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        text_encoder ([`T5EncoderModel`]):
-            Frozen text-encoder. Cosmos uses
-            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the
-            [t5-11b](https://huggingface.co/google-t5/t5-11b) variant.
-        tokenizer (`T5TokenizerFast`):
-            Tokenizer of class
-            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
-        transformer ([`CosmosTransformer3DModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`EDMEulerScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKLWan`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-    # We mark safety_checker as optional here to get around some test failures, but it is not really optional
-    _optional_components = ["safety_checker"]
-
-    def __init__(
-        self,
-        text_encoder: T5EncoderModel,
-        tokenizer: T5TokenizerFast,
-        transformer: CosmosTransformer3DModel,
-        vae: AutoencoderKLWan,
-        scheduler: EDMEulerScheduler,
-        safety_checker: CosmosSafetyChecker = None,
-    ):
-        super().__init__()
-
-        if safety_checker is None:
-            safety_checker = CosmosSafetyChecker()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-            safety_checker=safety_checker,
-        )
-
-        self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
-        self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        max_sequence_length: int = 512,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            return_tensors="pt",
-            return_length=True,
-            return_offsets_mapping=False,
-        )
-        text_input_ids = text_inputs.input_ids
-        prompt_attention_mask = text_inputs.attention_mask.bool().to(device)
-
-        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
-            logger.warning(
-                "The following part of your input was truncated because `max_sequence_length` is set to "
-                f" {max_sequence_length} tokens: {removed_text}"
-            )
-
-        prompt_embeds = self.text_encoder(
-            text_input_ids.to(device), attention_mask=prompt_attention_mask
-        ).last_hidden_state
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        lengths = prompt_attention_mask.sum(dim=1).cpu()
-        for i, length in enumerate(lengths):
-            prompt_embeds[i, length:] = 0
-
-        return prompt_embeds
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt with num_videos_per_prompt->num_images_per_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        do_classifier_free_guidance: bool = True,
-        num_images_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 512,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                Whether to use classifier free guidance or not.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            device: (`torch.device`, *optional*):
-                torch device
-            dtype: (`torch.dtype`, *optional*):
-                torch dtype
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
-            )
-
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            _, seq_len, _ = prompt_embeds.shape
-            prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-
-            negative_prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=negative_prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
-            )
-
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            _, seq_len, _ = negative_prompt_embeds.shape
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-        return prompt_embeds, negative_prompt_embeds
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: 16,
-        height: int = 768,
-        width: int = 1360,
-        num_frames: int = 1,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype) * self.scheduler.config.sigma_max
-
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latent_height = height // self.vae_scale_factor_spatial
-        latent_width = width // self.vae_scale_factor_spatial
-        shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents * self.scheduler.config.sigma_max
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-    ):
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def current_timestep(self):
-        return self._current_timestep
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        height: int = 768,
-        width: int = 1360,
-        num_inference_steps: int = 35,
-        guidance_scale: float = 7.0,
-        num_images_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback_on_step_end: Optional[
-            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
-        ] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, defaults to `768`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `1360`):
-                The width in pixels of the generated image.
-            num_inference_steps (`int`, defaults to `35`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            guidance_scale (`float`, defaults to `7.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion
-                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
-                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
-                `guidance_scale > 1`.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
-                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`CosmosImagePipelineOutput`] instead of a plain tuple.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-
-        Examples:
-
-        Returns:
-            [`~CosmosImagePipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`CosmosImagePipelineOutput`] is returned, otherwise a `tuple` is returned
-                where the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        if self.safety_checker is None:
-            raise ValueError(
-                f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
-                "[NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). "
-                f"Please ensure that you are compliant with the license agreement."
-            )
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        num_frames = 1
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs)
-
-        self._guidance_scale = guidance_scale
-        self._current_timestep = None
-        self._interrupt = False
-
-        device = self._execution_device
-
-        if self.safety_checker is not None:
-            self.safety_checker.to(device)
-            if prompt is not None:
-                prompt_list = [prompt] if isinstance(prompt, str) else prompt
-                for p in prompt_list:
-                    if not self.safety_checker.check_text_safety(p):
-                        raise ValueError(
-                            f"Cosmos Guardrail detected unsafe text in the prompt: {p}. Please ensure that the "
-                            f"prompt abides by the NVIDIA Open Model License Agreement."
-                        )
-            self.safety_checker.to("cpu")
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # 3. Encode input prompt
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            num_images_per_prompt=num_images_per_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            device=device,
-            max_sequence_length=max_sequence_length,
-        )
-
-        # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device)
-
-        # 5. Prepare latent variables
-        transformer_dtype = self.transformer.dtype
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-
-        padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
-
-        # 6. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                self._current_timestep = t
-                timestep = t.expand(latents.shape[0]).to(transformer_dtype)
-                current_sigma = self.scheduler.sigmas[i]
-
-                latent_model_input = latents
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-                latent_model_input = latent_model_input.to(transformer_dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    padding_mask=padding_mask,
-                    return_dict=False,
-                )[0]
-                noise_pred = self.scheduler.precondition_outputs(latents, noise_pred, current_sigma)
-
-                if self.do_classifier_free_guidance:
-                    noise_pred_uncond = self.transformer(
-                        hidden_states=latent_model_input,
-                        timestep=timestep,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        padding_mask=padding_mask,
-                        return_dict=False,
-                    )[0]
-                    noise_pred_uncond = self.scheduler.precondition_outputs(latents, noise_pred_uncond, current_sigma)
-                    noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond)
-
-                latents = self.scheduler.step(
-                    noise_pred, t, latents, pred_original_sample=noise_pred, return_dict=False
-                )[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        self._current_timestep = None
-
-        if not output_type == "latent":
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std / self.scheduler.config.sigma_data + latents_mean
-            video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
-
-            if self.safety_checker is not None:
-                self.safety_checker.to(device)
-                video = self.video_processor.postprocess_video(video, output_type="np")
-                video = (video * 255).astype(np.uint8)
-                video_batch = []
-                for vid in video:
-                    vid = self.safety_checker.check_video_safety(vid)
-                    video_batch.append(vid)
-                video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1
-                video = torch.from_numpy(video).permute(0, 4, 1, 2, 3)
-                video = self.video_processor.postprocess_video(video, output_type=output_type)
-                self.safety_checker.to("cpu")
-            else:
-                video = self.video_processor.postprocess_video(video, output_type=output_type)
-            image = [batch[0] for batch in video]
-            if isinstance(video, torch.Tensor):
-                image = torch.stack(image)
-            elif isinstance(video, np.ndarray):
-                image = np.stack(image)
-        else:
-            image = latents[:, :, 0]
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return CosmosImagePipelineOutput(images=image)

src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

@@ -1,771 +0,0 @@
-# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
-#
-# 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 inspect
-from typing import Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from transformers import T5EncoderModel, T5TokenizerFast
-
-from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...image_processor import PipelineImageInput
-from ...models import AutoencoderKLWan, CosmosTransformer3DModel
-from ...schedulers import EDMEulerScheduler
-from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring
-from ...utils.torch_utils import randn_tensor
-from ...video_processor import VideoProcessor
-from ..pipeline_utils import DiffusionPipeline
-from .pipeline_output import CosmosPipelineOutput
-
-
-if is_cosmos_guardrail_available():
-    from cosmos_guardrail import CosmosSafetyChecker
-else:
-
-    class CosmosSafetyChecker:
-        def __init__(self, *args, **kwargs):
-            raise ImportError(
-                "`cosmos_guardrail` is not installed. Please install it to use the safety checker for Cosmos: `pip install cosmos_guardrail`."
-            )
-
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```python
-        >>> import torch
-        >>> from diffusers import Cosmos2VideoToWorldPipeline
-        >>> from diffusers.utils import export_to_video, load_image
-
-        >>> # Available checkpoints: nvidia/Cosmos-Predict2-2B-Video2World, nvidia/Cosmos-Predict2-14B-Video2World
-        >>> model_id = "nvidia/Cosmos-Predict2-2B-Video2World"
-        >>> pipe = Cosmos2VideoToWorldPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-
-        >>> prompt = "A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
-        >>> negative_prompt = "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
-        >>> image = load_image(
-        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yellow-scrubber.png"
-        ... )
-
-        >>> video = pipe(
-        ...     image=image, prompt=prompt, negative_prompt=negative_prompt, generator=torch.Generator().manual_seed(1)
-        ... ).frames[0]
-        >>> export_to_video(video, "output.mp4", fps=16)
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
-def retrieve_latents(
-    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
-):
-    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
-        return encoder_output.latent_dist.sample(generator)
-    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
-        return encoder_output.latent_dist.mode()
-    elif hasattr(encoder_output, "latents"):
-        return encoder_output.latents
-    else:
-        raise AttributeError("Could not access latents of provided encoder_output")
-
-
-class Cosmos2VideoToWorldPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-image generation using [Cosmos](https://github.com/NVIDIA/Cosmos).
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        text_encoder ([`T5EncoderModel`]):
-            Frozen text-encoder. Cosmos uses
-            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel); specifically the
-            [t5-11b](https://huggingface.co/google-t5/t5-11b) variant.
-        tokenizer (`T5TokenizerFast`):
-            Tokenizer of class
-            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
-        transformer ([`CosmosTransformer3DModel`]):
-            Conditional Transformer to denoise the encoded image latents.
-        scheduler ([`EDMEulerScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKLWan`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-    # We mark safety_checker as optional here to get around some test failures, but it is not really optional
-    _optional_components = ["safety_checker"]
-
-    def __init__(
-        self,
-        text_encoder: T5EncoderModel,
-        tokenizer: T5TokenizerFast,
-        transformer: CosmosTransformer3DModel,
-        vae: AutoencoderKLWan,
-        scheduler: EDMEulerScheduler,
-        safety_checker: CosmosSafetyChecker = None,
-    ):
-        super().__init__()
-
-        if safety_checker is None:
-            safety_checker = CosmosSafetyChecker()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-            safety_checker=safety_checker,
-        )
-
-        self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
-        self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline._get_t5_prompt_embeds
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        max_sequence_length: int = 512,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            return_tensors="pt",
-            return_length=True,
-            return_offsets_mapping=False,
-        )
-        text_input_ids = text_inputs.input_ids
-        prompt_attention_mask = text_inputs.attention_mask.bool().to(device)
-
-        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
-            logger.warning(
-                "The following part of your input was truncated because `max_sequence_length` is set to "
-                f" {max_sequence_length} tokens: {removed_text}"
-            )
-
-        prompt_embeds = self.text_encoder(
-            text_input_ids.to(device), attention_mask=prompt_attention_mask
-        ).last_hidden_state
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        lengths = prompt_attention_mask.sum(dim=1).cpu()
-        for i, length in enumerate(lengths):
-            prompt_embeds[i, length:] = 0
-
-        return prompt_embeds
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        do_classifier_free_guidance: bool = True,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 512,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                Whether to use classifier free guidance or not.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            device: (`torch.device`, *optional*):
-                torch device
-            dtype: (`torch.dtype`, *optional*):
-                torch dtype
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
-            )
-
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            _, seq_len, _ = prompt_embeds.shape
-            prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-            prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-
-            negative_prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=negative_prompt, max_sequence_length=max_sequence_length, device=device, dtype=dtype
-            )
-
-            # duplicate text embeddings for each generation per prompt, using mps friendly method
-            _, seq_len, _ = negative_prompt_embeds.shape
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        return prompt_embeds, negative_prompt_embeds
-
-    def prepare_latents(
-        self,
-        video: torch.Tensor,
-        batch_size: int,
-        num_channels_latents: 16,
-        height: int = 704,
-        width: int = 1280,
-        num_frames: int = 77,
-        do_classifier_free_guidance: bool = True,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        num_cond_frames = video.size(2)
-        if num_cond_frames >= num_frames:
-            # Take the last `num_frames` frames for conditioning
-            num_cond_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-            video = video[:, :, -num_frames:]
-        else:
-            num_cond_latent_frames = (num_cond_frames - 1) // self.vae_scale_factor_temporal + 1
-            num_padding_frames = num_frames - num_cond_frames
-            last_frame = video[:, :, -1:]
-            padding = last_frame.repeat(1, 1, num_padding_frames, 1, 1)
-            video = torch.cat([video, padding], dim=2)
-
-        if isinstance(generator, list):
-            init_latents = [
-                retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator=generator[i])
-                for i in range(batch_size)
-            ]
-        else:
-            init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
-
-        init_latents = torch.cat(init_latents, dim=0).to(dtype)
-
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(device, dtype)
-        )
-        init_latents = (init_latents - latents_mean) / latents_std * self.scheduler.config.sigma_data
-
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latent_height = height // self.vae_scale_factor_spatial
-        latent_width = width // self.vae_scale_factor_spatial
-        shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device=device, dtype=dtype)
-
-        latents = latents * self.scheduler.config.sigma_max
-
-        padding_shape = (batch_size, 1, num_latent_frames, latent_height, latent_width)
-        ones_padding = latents.new_ones(padding_shape)
-        zeros_padding = latents.new_zeros(padding_shape)
-
-        cond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
-        cond_indicator[:, :, :num_cond_latent_frames] = 1.0
-        cond_mask = cond_indicator * ones_padding + (1 - cond_indicator) * zeros_padding
-
-        uncond_indicator = uncond_mask = None
-        if do_classifier_free_guidance:
-            uncond_indicator = latents.new_zeros(1, 1, latents.size(2), 1, 1)
-            uncond_indicator[:, :, :num_cond_latent_frames] = 1.0
-            uncond_mask = uncond_indicator * ones_padding + (1 - uncond_indicator) * zeros_padding
-
-        return latents, init_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask
-
-    # Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-    ):
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def current_timestep(self):
-        return self._current_timestep
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        image: PipelineImageInput = None,
-        video: List[PipelineImageInput] = None,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        height: int = 704,
-        width: int = 1280,
-        num_frames: int = 77,
-        num_inference_steps: int = 35,
-        guidance_scale: float = 7.0,
-        fps: int = 16,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback_on_step_end: Optional[
-            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
-        ] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 512,
-        sigma_conditioning: float = 0.0001,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*):
-                The image to be used as a conditioning input for the video generation.
-            video (`List[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
-                The video to be used as a conditioning input for the video generation.
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, defaults to `704`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `1280`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `77`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `35`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            guidance_scale (`float`, defaults to `7.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion
-                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
-                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
-                `guidance_scale > 1`.
-            fps (`int`, defaults to `16`):
-                The frames per second of the generated video.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
-                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`CosmosPipelineOutput`] instead of a plain tuple.
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int`, defaults to `512`):
-                The maximum number of tokens in the prompt. If the prompt exceeds this length, it will be truncated. If
-                the prompt is shorter than this length, it will be padded.
-            sigma_conditioning (`float`, defaults to `0.0001`):
-                The sigma value used for scaling conditioning latents. Ideally, it should not be changed or should be
-                set to a small value close to zero.
-
-        Examples:
-
-        Returns:
-            [`~CosmosPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`CosmosPipelineOutput`] is returned, otherwise a `tuple` is returned where
-                the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        if self.safety_checker is None:
-            raise ValueError(
-                f"You have disabled the safety checker for {self.__class__}. This is in violation of the "
-                "[NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). "
-                f"Please ensure that you are compliant with the license agreement."
-            )
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(prompt, height, width, prompt_embeds, callback_on_step_end_tensor_inputs)
-
-        self._guidance_scale = guidance_scale
-        self._current_timestep = None
-        self._interrupt = False
-
-        device = self._execution_device
-
-        if self.safety_checker is not None:
-            self.safety_checker.to(device)
-            if prompt is not None:
-                prompt_list = [prompt] if isinstance(prompt, str) else prompt
-                for p in prompt_list:
-                    if not self.safety_checker.check_text_safety(p):
-                        raise ValueError(
-                            f"Cosmos Guardrail detected unsafe text in the prompt: {p}. Please ensure that the "
-                            f"prompt abides by the NVIDIA Open Model License Agreement."
-                        )
-            self.safety_checker.to("cpu")
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # 3. Encode input prompt
-        (
-            prompt_embeds,
-            negative_prompt_embeds,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            device=device,
-            max_sequence_length=max_sequence_length,
-        )
-
-        # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device)
-
-        # 5. Prepare latent variables
-        vae_dtype = self.vae.dtype
-        transformer_dtype = self.transformer.dtype
-
-        if image is not None:
-            video = self.video_processor.preprocess(image, height, width).unsqueeze(2)
-        else:
-            video = self.video_processor.preprocess_video(video, height, width)
-        video = video.to(device=device, dtype=vae_dtype)
-
-        num_channels_latents = self.transformer.config.in_channels - 1
-        latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = self.prepare_latents(
-            video,
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            self.do_classifier_free_guidance,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-        unconditioning_latents = None
-
-        cond_mask = cond_mask.to(transformer_dtype)
-        if self.do_classifier_free_guidance:
-            uncond_mask = uncond_mask.to(transformer_dtype)
-            unconditioning_latents = conditioning_latents
-
-        padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
-        sigma_conditioning = torch.tensor(sigma_conditioning, dtype=torch.float32, device=device)
-        t_conditioning = self.scheduler.precondition_noise(sigma_conditioning)
-
-        # 6. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                self._current_timestep = t
-                timestep = t.view(1, 1, 1, 1, 1).expand(
-                    latents.size(0), -1, latents.size(2), -1, -1
-                )  # [B, 1, T, 1, 1]
-                current_sigma = self.scheduler.sigmas[i]
-
-                cond_latent = self.scheduler.scale_model_input(latents, t)
-                cond_latent = cond_indicator * conditioning_latents + (1 - cond_indicator) * cond_latent
-                cond_latent = cond_latent.to(transformer_dtype)
-                cond_timestep = cond_indicator * t_conditioning + (1 - cond_indicator) * timestep
-                cond_timestep = cond_timestep.to(transformer_dtype)
-
-                noise_pred = self.transformer(
-                    hidden_states=cond_latent,
-                    timestep=cond_timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    fps=fps,
-                    condition_mask=cond_mask,
-                    padding_mask=padding_mask,
-                    return_dict=False,
-                )[0]
-                noise_pred = self.scheduler.precondition_outputs(latents, noise_pred, current_sigma)
-                noise_pred = cond_indicator * conditioning_latents + (1 - cond_indicator) * noise_pred
-
-                if self.do_classifier_free_guidance:
-                    uncond_latent = self.scheduler.scale_model_input(latents, t)
-                    uncond_latent = uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * uncond_latent
-                    uncond_latent = uncond_latent.to(transformer_dtype)
-                    uncond_timestep = uncond_indicator * t_conditioning + (1 - uncond_indicator) * timestep
-                    uncond_timestep = uncond_timestep.to(transformer_dtype)
-
-                    noise_pred_uncond = self.transformer(
-                        hidden_states=uncond_latent,
-                        timestep=uncond_timestep,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        fps=fps,
-                        condition_mask=uncond_mask,
-                        padding_mask=padding_mask,
-                        return_dict=False,
-                    )[0]
-                    noise_pred_uncond = self.scheduler.precondition_outputs(latents, noise_pred_uncond, current_sigma)
-                    noise_pred_uncond = (
-                        uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * noise_pred_uncond
-                    )
-                    noise_pred = noise_pred + self.guidance_scale * (noise_pred - noise_pred_uncond)
-
-                latents = self.scheduler.step(
-                    noise_pred, t, latents, pred_original_sample=noise_pred, return_dict=False
-                )[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        self._current_timestep = None
-
-        if not output_type == "latent":
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = (
-                torch.tensor(self.vae.config.latents_std)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents = latents * latents_std / self.scheduler.config.sigma_data + latents_mean
-            video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
-
-            if self.safety_checker is not None:
-                self.safety_checker.to(device)
-                video = self.video_processor.postprocess_video(video, output_type="np")
-                video = (video * 255).astype(np.uint8)
-                video_batch = []
-                for vid in video:
-                    vid = self.safety_checker.check_video_safety(vid)
-                    video_batch.append(vid)
-                video = np.stack(video_batch).astype(np.float32) / 255.0 * 2 - 1
-                video = torch.from_numpy(video).permute(0, 4, 1, 2, 3)
-                video = self.video_processor.postprocess_video(video, output_type=output_type)
-                self.safety_checker.to("cpu")
-            else:
-                video = self.video_processor.postprocess_video(video, output_type=output_type)
-        else:
-            video = latents
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (video,)
-
-        return CosmosPipelineOutput(frames=video)

src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py

@@ -426,12 +426,12 @@ class CosmosTextToWorldPipeline(DiffusionPipeline):
                 The height in pixels of the generated image.
             width (`int`, defaults to `1280`):
                 The width in pixels of the generated image.
-            num_frames (`int`, defaults to `121`):
+            num_frames (`int`, defaults to `129`):
                 The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `36`):
+            num_inference_steps (`int`, defaults to `50`):
                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                 expense of slower inference.
-            guidance_scale (`float`, defaults to `7.0`):
+            guidance_scale (`float`, defaults to `6.0`):
                 Guidance scale as defined in [Classifier-Free Diffusion
                 Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                 of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
@@ -457,6 +457,9 @@ class CosmosTextToWorldPipeline(DiffusionPipeline):
                 The output format of the generated image. Choose between `PIL.Image` or `np.array`.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`CosmosPipelineOutput`] instead of a plain tuple.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
             callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
                 A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
                 each denoising step during the inference. with the following arguments: `callback_on_step_end(self:

src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

@@ -541,12 +541,12 @@ class CosmosVideoToWorldPipeline(DiffusionPipeline):
                 The height in pixels of the generated image.
             width (`int`, defaults to `1280`):
                 The width in pixels of the generated image.
-            num_frames (`int`, defaults to `121`):
+            num_frames (`int`, defaults to `129`):
                 The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `36`):
+            num_inference_steps (`int`, defaults to `50`):
                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                 expense of slower inference.
-            guidance_scale (`float`, defaults to `7.0`):
+            guidance_scale (`float`, defaults to `6.0`):
                 Guidance scale as defined in [Classifier-Free Diffusion
                 Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
                 of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
@@ -572,6 +572,9 @@ class CosmosVideoToWorldPipeline(DiffusionPipeline):
                 The output format of the generated image. Choose between `PIL.Image` or `np.array`.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`CosmosPipelineOutput`] instead of a plain tuple.
+            clip_skip (`int`, *optional*):
+                Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+                the output of the pre-final layer will be used for computing the prompt embeddings.
             callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
                 A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
                 each denoising step during the inference. with the following arguments: `callback_on_step_end(self:

src/diffusers/pipelines/cosmos/pipeline_output.py

@@ -1,20 +1,14 @@
 from dataclasses import dataclass
-from typing import List, Union
 
-import numpy as np
-import PIL.Image
 import torch
 
-from diffusers.utils import BaseOutput, get_logger
-
-
-logger = get_logger(__name__)
+from diffusers.utils import BaseOutput
 
 
 @dataclass
 class CosmosPipelineOutput(BaseOutput):
     r"""
-    Output class for Cosmos any-to-world/video pipelines.
+    Output class for Cosmos pipelines.
 
     Args:
         frames (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
@@ -24,17 +18,3 @@ class CosmosPipelineOutput(BaseOutput):
     """
 
     frames: torch.Tensor
-
-
-@dataclass
-class CosmosImagePipelineOutput(BaseOutput):
-    """
-    Output class for CogView3 pipelines.
-
-    Args:
-        images (`List[PIL.Image.Image]` or `np.ndarray`)
-            List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
-            num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
-    """
-
-    images: Union[List[PIL.Image.Image], np.ndarray]

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -700,7 +700,7 @@ class FluxPipeline(
             latents (`torch.FloatTensor`, *optional*):
                 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
                 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
+                tensor will ge generated by sampling using the supplied random `generator`.
             prompt_embeds (`torch.FloatTensor`, *optional*):
                 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
                 provided, text embeddings will be generated from `prompt` input argument.

src/diffusers/pipelines/sana/__init__.py

@@ -25,7 +25,6 @@ else:
     _import_structure["pipeline_sana"] = ["SanaPipeline"]
     _import_structure["pipeline_sana_controlnet"] = ["SanaControlNetPipeline"]
     _import_structure["pipeline_sana_sprint"] = ["SanaSprintPipeline"]
-    _import_structure["pipeline_sana_sprint_img2img"] = ["SanaSprintImg2ImgPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -38,7 +37,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
         from .pipeline_sana import SanaPipeline
         from .pipeline_sana_controlnet import SanaControlNetPipeline
         from .pipeline_sana_sprint import SanaSprintPipeline
-        from .pipeline_sana_sprint_img2img import SanaSprintImg2ImgPipeline
 else:
     import sys
 

src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py

@@ -1,975 +0,0 @@
-# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
-#
-# 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 html
-import inspect
-import re
-import urllib.parse as ul
-import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-from transformers import Gemma2PreTrainedModel, GemmaTokenizer, GemmaTokenizerFast
-
-from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...image_processor import PipelineImageInput, PixArtImageProcessor
-from ...loaders import SanaLoraLoaderMixin
-from ...models import AutoencoderDC, SanaTransformer2DModel
-from ...schedulers import DPMSolverMultistepScheduler
-from ...utils import (
-    BACKENDS_MAPPING,
-    USE_PEFT_BACKEND,
-    is_bs4_available,
-    is_ftfy_available,
-    is_torch_xla_available,
-    logging,
-    replace_example_docstring,
-    scale_lora_layers,
-    unscale_lora_layers,
-)
-from ...utils.torch_utils import randn_tensor
-from ..pipeline_utils import DiffusionPipeline
-from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN
-from .pipeline_output import SanaPipelineOutput
-
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-if is_bs4_available():
-    from bs4 import BeautifulSoup
-
-if is_ftfy_available():
-    import ftfy
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import SanaSprintImg2ImgPipeline
-        >>> from diffusers.utils.loading_utils import load_image
-
-        >>> pipe = SanaSprintImg2ImgPipeline.from_pretrained(
-        ...     "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", torch_dtype=torch.bfloat16
-        ... )
-        >>> pipe.to("cuda")
-
-        >>> image = load_image(
-        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
-        ... )
-
-
-        >>> image = pipe(prompt="a cute pink bear", image=image, strength=0.5, height=832, width=480).images[0]
-        >>> image[0].save("output.png")
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class SanaSprintImg2ImgPipeline(DiffusionPipeline, SanaLoraLoaderMixin):
-    r"""
-    Pipeline for text-to-image generation using [SANA-Sprint](https://huggingface.co/papers/2503.09641).
-    """
-
-    # fmt: off
-    bad_punct_regex = re.compile(
-        r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + "\\" + r"\/" + r"\*" + r"]{1,}")
-    # fmt: on
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds"]
-
-    def __init__(
-        self,
-        tokenizer: Union[GemmaTokenizer, GemmaTokenizerFast],
-        text_encoder: Gemma2PreTrainedModel,
-        vae: AutoencoderDC,
-        transformer: SanaTransformer2DModel,
-        scheduler: DPMSolverMultistepScheduler,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
-        )
-
-        self.vae_scale_factor = (
-            2 ** (len(self.vae.config.encoder_block_out_channels) - 1)
-            if hasattr(self, "vae") and self.vae is not None
-            else 32
-        )
-        self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
-
-    # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_slicing
-    def enable_vae_slicing(self):
-        r"""
-        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
-        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
-        """
-        self.vae.enable_slicing()
-
-    # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.disable_vae_slicing
-    def disable_vae_slicing(self):
-        r"""
-        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_slicing()
-
-    # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_tiling
-    def enable_vae_tiling(self):
-        r"""
-        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
-        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
-        processing larger images.
-        """
-        self.vae.enable_tiling()
-
-    def disable_vae_tiling(self):
-        r"""
-        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
-        computing decoding in one step.
-        """
-        self.vae.disable_tiling()
-
-    # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds
-    def _get_gemma_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]],
-        device: torch.device,
-        dtype: torch.dtype,
-        clean_caption: bool = False,
-        max_sequence_length: int = 300,
-        complex_human_instruction: Optional[List[str]] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            device: (`torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            clean_caption (`bool`, defaults to `False`):
-                If `True`, the function will preprocess and clean the provided caption before encoding.
-            max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt.
-            complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`):
-                If `complex_human_instruction` is not empty, the function will use the complex Human instruction for
-                the prompt.
-        """
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-
-        if getattr(self, "tokenizer", None) is not None:
-            self.tokenizer.padding_side = "right"
-
-        prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
-
-        # prepare complex human instruction
-        if not complex_human_instruction:
-            max_length_all = max_sequence_length
-        else:
-            chi_prompt = "\n".join(complex_human_instruction)
-            prompt = [chi_prompt + p for p in prompt]
-            num_chi_prompt_tokens = len(self.tokenizer.encode(chi_prompt))
-            max_length_all = num_chi_prompt_tokens + max_sequence_length - 2
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_length_all,
-            truncation=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-
-        prompt_attention_mask = text_inputs.attention_mask
-        prompt_attention_mask = prompt_attention_mask.to(device)
-
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
-        prompt_embeds = prompt_embeds[0].to(dtype=dtype, device=device)
-
-        return prompt_embeds, prompt_attention_mask
-
-    # Copied from diffusers.pipelines.sana.pipeline_sana_sprint.SanaSprintPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        num_images_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        clean_caption: bool = False,
-        max_sequence_length: int = 300,
-        complex_human_instruction: Optional[List[str]] = None,
-        lora_scale: Optional[float] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                number of images that should be generated per prompt
-            device: (`torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            clean_caption (`bool`, defaults to `False`):
-                If `True`, the function will preprocess and clean the provided caption before encoding.
-            max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt.
-            complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`):
-                If `complex_human_instruction` is not empty, the function will use the complex Human instruction for
-                the prompt.
-        """
-
-        if device is None:
-            device = self._execution_device
-
-        if self.text_encoder is not None:
-            dtype = self.text_encoder.dtype
-        else:
-            dtype = None
-
-        # set lora scale so that monkey patched LoRA
-        # function of text encoder can correctly access it
-        if lora_scale is not None and isinstance(self, SanaLoraLoaderMixin):
-            self._lora_scale = lora_scale
-
-            # dynamically adjust the LoRA scale
-            if self.text_encoder is not None and USE_PEFT_BACKEND:
-                scale_lora_layers(self.text_encoder, lora_scale)
-
-        if getattr(self, "tokenizer", None) is not None:
-            self.tokenizer.padding_side = "right"
-
-        # See Section 3.1. of the paper.
-        max_length = max_sequence_length
-        select_index = [0] + list(range(-max_length + 1, 0))
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
-                prompt=prompt,
-                device=device,
-                dtype=dtype,
-                clean_caption=clean_caption,
-                max_sequence_length=max_sequence_length,
-                complex_human_instruction=complex_human_instruction,
-            )
-
-            prompt_embeds = prompt_embeds[:, select_index]
-            prompt_attention_mask = prompt_attention_mask[:, select_index]
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
-
-        if self.text_encoder is not None:
-            if isinstance(self, SanaLoraLoaderMixin) and USE_PEFT_BACKEND:
-                # Retrieve the original scale by scaling back the LoRA layers
-                unscale_lora_layers(self.text_encoder, lora_scale)
-
-        return prompt_embeds, prompt_attention_mask
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://huggingface.co/papers/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps
-    def get_timesteps(self, num_inference_steps, strength, device):
-        # get the original timestep using init_timestep
-        init_timestep = min(num_inference_steps * strength, num_inference_steps)
-
-        t_start = int(max(num_inference_steps - init_timestep, 0))
-        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
-        if hasattr(self.scheduler, "set_begin_index"):
-            self.scheduler.set_begin_index(t_start * self.scheduler.order)
-
-        return timesteps, num_inference_steps - t_start
-
-    def check_inputs(
-        self,
-        prompt,
-        strength,
-        height,
-        width,
-        num_inference_steps,
-        timesteps,
-        max_timesteps,
-        intermediate_timesteps,
-        callback_on_step_end_tensor_inputs=None,
-        prompt_embeds=None,
-        prompt_attention_mask=None,
-    ):
-        if strength < 0 or strength > 1:
-            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
-
-        if height % 32 != 0 or width % 32 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if prompt_embeds is not None and prompt_attention_mask is None:
-            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
-
-        if timesteps is not None and len(timesteps) != num_inference_steps + 1:
-            raise ValueError("If providing custom timesteps, `timesteps` must be of length `num_inference_steps + 1`.")
-
-        if timesteps is not None and max_timesteps is not None:
-            raise ValueError("If providing custom timesteps, `max_timesteps` should not be provided.")
-
-        if timesteps is None and max_timesteps is None:
-            raise ValueError("Should provide either `timesteps` or `max_timesteps`.")
-
-        if intermediate_timesteps is not None and num_inference_steps != 2:
-            raise ValueError("Intermediate timesteps for SCM is not supported when num_inference_steps != 2.")
-
-    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
-    def _text_preprocessing(self, text, clean_caption=False):
-        if clean_caption and not is_bs4_available():
-            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
-            logger.warning("Setting `clean_caption` to False...")
-            clean_caption = False
-
-        if clean_caption and not is_ftfy_available():
-            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
-            logger.warning("Setting `clean_caption` to False...")
-            clean_caption = False
-
-        if not isinstance(text, (tuple, list)):
-            text = [text]
-
-        def process(text: str):
-            if clean_caption:
-                text = self._clean_caption(text)
-                text = self._clean_caption(text)
-            else:
-                text = text.lower().strip()
-            return text
-
-        return [process(t) for t in text]
-
-    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
-    def _clean_caption(self, caption):
-        caption = str(caption)
-        caption = ul.unquote_plus(caption)
-        caption = caption.strip().lower()
-        caption = re.sub("<person>", "person", caption)
-        # urls:
-        caption = re.sub(
-            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
-            "",
-            caption,
-        )  # regex for urls
-        caption = re.sub(
-            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
-            "",
-            caption,
-        )  # regex for urls
-        # html:
-        caption = BeautifulSoup(caption, features="html.parser").text
-
-        # @<nickname>
-        caption = re.sub(r"@[\w\d]+\b", "", caption)
-
-        # 31C0—31EF CJK Strokes
-        # 31F0—31FF Katakana Phonetic Extensions
-        # 3200—32FF Enclosed CJK Letters and Months
-        # 3300—33FF CJK Compatibility
-        # 3400—4DBF CJK Unified Ideographs Extension A
-        # 4DC0—4DFF Yijing Hexagram Symbols
-        # 4E00—9FFF CJK Unified Ideographs
-        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
-        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
-        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
-        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
-        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
-        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
-        caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
-        #######################################################
-
-        # все виды тире / all types of dash --> "-"
-        caption = re.sub(
-            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
-            "-",
-            caption,
-        )
-
-        # кавычки к одному стандарту
-        caption = re.sub(r"[`´«»“”¨]", '"', caption)
-        caption = re.sub(r"[‘’]", "'", caption)
-
-        # &quot;
-        caption = re.sub(r"&quot;?", "", caption)
-        # &amp
-        caption = re.sub(r"&amp", "", caption)
-
-        # ip addresses:
-        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
-
-        # article ids:
-        caption = re.sub(r"\d:\d\d\s+$", "", caption)
-
-        # \n
-        caption = re.sub(r"\\n", " ", caption)
-
-        # "#123"
-        caption = re.sub(r"#\d{1,3}\b", "", caption)
-        # "#12345.."
-        caption = re.sub(r"#\d{5,}\b", "", caption)
-        # "123456.."
-        caption = re.sub(r"\b\d{6,}\b", "", caption)
-        # filenames:
-        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
-
-        #
-        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
-        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
-
-        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
-        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
-
-        # this-is-my-cute-cat / this_is_my_cute_cat
-        regex2 = re.compile(r"(?:\-|\_)")
-        if len(re.findall(regex2, caption)) > 3:
-            caption = re.sub(regex2, " ", caption)
-
-        caption = ftfy.fix_text(caption)
-        caption = html.unescape(html.unescape(caption))
-
-        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
-        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
-        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
-
-        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
-        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
-        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
-        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
-        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
-
-        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
-
-        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
-
-        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
-        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
-        caption = re.sub(r"\s+", " ", caption)
-
-        caption.strip()
-
-        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
-        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
-        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
-        caption = re.sub(r"^\.\S+$", "", caption)
-
-        return caption.strip()
-
-    def prepare_image(
-        self,
-        image: PipelineImageInput,
-        width: int,
-        height: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ):
-        if isinstance(image, torch.Tensor):
-            if image.ndim == 3:
-                image = image.unsqueeze(0)
-            # Resize if current dimensions do not match target dimensions.
-            if image.shape[2] != height or image.shape[3] != width:
-                image = F.interpolate(image, size=(height, width), mode="bilinear", align_corners=False)
-
-            image = self.image_processor.preprocess(image, height=height, width=width)
-
-        else:
-            image = self.image_processor.preprocess(image, height=height, width=width)
-
-        image = image.to(device=device, dtype=dtype)
-
-        return image
-
-    def prepare_latents(
-        self, image, timestep, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None
-    ):
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        shape = (
-            batch_size,
-            num_channels_latents,
-            int(height) // self.vae_scale_factor,
-            int(width) // self.vae_scale_factor,
-        )
-
-        if image.shape[1] != num_channels_latents:
-            image = self.vae.encode(image).latent
-            image_latents = image * self.vae.config.scaling_factor * self.scheduler.config.sigma_data
-        else:
-            image_latents = image
-        if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
-            # expand init_latents for batch_size
-            additional_image_per_prompt = batch_size // image_latents.shape[0]
-            image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
-        elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
-            raise ValueError(
-                f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
-            )
-        else:
-            image_latents = torch.cat([image_latents], dim=0)
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        # adapt from https://github.com/huggingface/diffusers/blob/c36f8487df35895421c15f351c7d360bd680[…]/examples/research_projects/sana/train_sana_sprint_diffusers.py
-        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) * self.scheduler.config.sigma_data
-        latents = torch.cos(timestep) * image_latents + torch.sin(timestep) * noise
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        num_inference_steps: int = 2,
-        timesteps: List[int] = None,
-        max_timesteps: float = 1.57080,
-        intermediate_timesteps: float = 1.3,
-        guidance_scale: float = 4.5,
-        image: PipelineImageInput = None,
-        strength: float = 0.6,
-        num_images_per_prompt: Optional[int] = 1,
-        height: int = 1024,
-        width: int = 1024,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_attention_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        clean_caption: bool = False,
-        use_resolution_binning: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 300,
-        complex_human_instruction: List[str] = [
-            "Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:",
-            "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.",
-            "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.",
-            "Here are examples of how to transform or refine prompts:",
-            "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.",
-            "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.",
-            "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:",
-            "User Prompt: ",
-        ],
-    ) -> Union[SanaPipelineOutput, Tuple]:
-        """
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            num_inference_steps (`int`, *optional*, defaults to 20):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            max_timesteps (`float`, *optional*, defaults to 1.57080):
-                The maximum timestep value used in the SCM scheduler.
-            intermediate_timesteps (`float`, *optional*, defaults to 1.3):
-                The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2).
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
-                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
-                passed will be used. Must be in descending order.
-            guidance_scale (`float`, *optional*, defaults to 4.5):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size):
-                The width in pixels of the generated image.
-            eta (`float`, *optional*, defaults to 0.0):
-                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
-                [`schedulers.DDIMScheduler`], will be ignored for others.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will ge generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
-            attention_kwargs:
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            clean_caption (`bool`, *optional*, defaults to `True`):
-                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
-                be installed. If the dependencies are not installed, the embeddings will be created from the raw
-                prompt.
-            use_resolution_binning (`bool` defaults to `True`):
-                If set to `True`, the requested height and width are first mapped to the closest resolutions using
-                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
-                the requested resolution. Useful for generating non-square images.
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to `300`):
-                Maximum sequence length to use with the `prompt`.
-            complex_human_instruction (`List[str]`, *optional*):
-                Instructions for complex human attention:
-                https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.sana.pipeline_output.SanaPipelineOutput`] is returned,
-                otherwise a `tuple` is returned where the first element is a list with the generated images
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # 1. Check inputs. Raise error if not correct
-        if use_resolution_binning:
-            if self.transformer.config.sample_size == 32:
-                aspect_ratio_bin = ASPECT_RATIO_1024_BIN
-            else:
-                raise ValueError("Invalid sample size")
-            orig_height, orig_width = height, width
-            height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
-
-        self.check_inputs(
-            prompt=prompt,
-            strength=strength,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            timesteps=timesteps,
-            max_timesteps=max_timesteps,
-            intermediate_timesteps=intermediate_timesteps,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            prompt_embeds=prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-
-        # 2. Default height and width to transformer
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-        lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None
-
-        # 2. Preprocess image
-        init_image = self.prepare_image(image, width, height, device, self.vae.dtype)
-
-        # 3. Encode input prompt
-        (
-            prompt_embeds,
-            prompt_attention_mask,
-        ) = self.encode_prompt(
-            prompt,
-            num_images_per_prompt=num_images_per_prompt,
-            device=device,
-            prompt_embeds=prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            clean_caption=clean_caption,
-            max_sequence_length=max_sequence_length,
-            complex_human_instruction=complex_human_instruction,
-            lora_scale=lora_scale,
-        )
-
-        # 5. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            timesteps,
-            sigmas=None,
-            max_timesteps=max_timesteps,
-            intermediate_timesteps=intermediate_timesteps,
-        )
-        if hasattr(self.scheduler, "set_begin_index"):
-            self.scheduler.set_begin_index(0)
-
-        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
-        if num_inference_steps < 1:
-            raise ValueError(
-                f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline"
-                f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline."
-            )
-        latent_timestep = timesteps[:1]
-
-        # 5. Prepare latents.
-        latent_channels = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            init_image,
-            latent_timestep,
-            batch_size * num_images_per_prompt,
-            latent_channels,
-            height,
-            width,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-
-        # I think this is redundant given the scaling in prepare_latents
-        # latents = latents * self.scheduler.config.sigma_data
-
-        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
-        guidance = guidance.expand(latents.shape[0]).to(prompt_embeds.dtype)
-        guidance = guidance * self.transformer.config.guidance_embeds_scale
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 7. Denoising loop
-        timesteps = timesteps[:-1]
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        transformer_dtype = self.transformer.dtype
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latents.shape[0])
-                latents_model_input = latents / self.scheduler.config.sigma_data
-
-                scm_timestep = torch.sin(timestep) / (torch.cos(timestep) + torch.sin(timestep))
-
-                scm_timestep_expanded = scm_timestep.view(-1, 1, 1, 1)
-                latent_model_input = latents_model_input * torch.sqrt(
-                    scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2
-                )
-
-                # predict noise model_output
-                noise_pred = self.transformer(
-                    latent_model_input.to(dtype=transformer_dtype),
-                    encoder_hidden_states=prompt_embeds.to(dtype=transformer_dtype),
-                    encoder_attention_mask=prompt_attention_mask,
-                    guidance=guidance,
-                    timestep=scm_timestep,
-                    return_dict=False,
-                    attention_kwargs=self.attention_kwargs,
-                )[0]
-
-                noise_pred = (
-                    (1 - 2 * scm_timestep_expanded) * latent_model_input
-                    + (1 - 2 * scm_timestep_expanded + 2 * scm_timestep_expanded**2) * noise_pred
-                ) / torch.sqrt(scm_timestep_expanded**2 + (1 - scm_timestep_expanded) ** 2)
-                noise_pred = noise_pred.float() * self.scheduler.config.sigma_data
-
-                # compute previous image: x_t -> x_t-1
-                latents, denoised = self.scheduler.step(
-                    noise_pred, timestep, latents, **extra_step_kwargs, return_dict=False
-                )
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        latents = denoised / self.scheduler.config.sigma_data
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = latents.to(self.vae.dtype)
-            try:
-                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except torch.cuda.OutOfMemoryError as e:
-                warnings.warn(
-                    f"{e}. \n"
-                    f"Try to use VAE tiling for large images. For example: \n"
-                    f"pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)"
-                )
-            if use_resolution_binning:
-                image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
-
-        if not output_type == "latent":
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return SanaPipelineOutput(images=image)

src/diffusers/pipelines/wan/__init__.py

@@ -24,7 +24,6 @@ except OptionalDependencyNotAvailable:
 else:
     _import_structure["pipeline_wan"] = ["WanPipeline"]
     _import_structure["pipeline_wan_i2v"] = ["WanImageToVideoPipeline"]
-    _import_structure["pipeline_wan_vace"] = ["WanVACEPipeline"]
     _import_structure["pipeline_wan_video2video"] = ["WanVideoToVideoPipeline"]
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -36,7 +35,6 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     else:
         from .pipeline_wan import WanPipeline
         from .pipeline_wan_i2v import WanImageToVideoPipeline
-        from .pipeline_wan_vace import WanVACEPipeline
         from .pipeline_wan_video2video import WanVideoToVideoPipeline
 
 else:

src/diffusers/pipelines/wan/pipeline_wan_vace.py

@@ -1,950 +0,0 @@
-# Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
-#
-# 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 html
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import PIL.Image
-import regex as re
-import torch
-from transformers import AutoTokenizer, UMT5EncoderModel
-
-from ...callbacks import MultiPipelineCallbacks, PipelineCallback
-from ...image_processor import PipelineImageInput
-from ...loaders import WanLoraLoaderMixin
-from ...models import AutoencoderKLWan, WanVACETransformer3DModel
-from ...schedulers import FlowMatchEulerDiscreteScheduler
-from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring
-from ...utils.torch_utils import randn_tensor
-from ...video_processor import VideoProcessor
-from ..pipeline_utils import DiffusionPipeline
-from .pipeline_output import WanPipelineOutput
-
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-if is_ftfy_available():
-    import ftfy
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```python
-        >>> import torch
-        >>> import PIL.Image
-        >>> from diffusers import AutoencoderKLWan, WanVACEPipeline
-        >>> from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
-        >>> from diffusers.utils import export_to_video, load_image
-        def prepare_video_and_mask(first_img: PIL.Image.Image, last_img: PIL.Image.Image, height: int, width: int, num_frames: int):
-            first_img = first_img.resize((width, height))
-            last_img = last_img.resize((width, height))
-            frames = []
-            frames.append(first_img)
-            # Ideally, this should be 127.5 to match original code, but they perform computation on numpy arrays
-            # whereas we are passing PIL images. If you choose to pass numpy arrays, you can set it to 127.5 to
-            # match the original code.
-            frames.extend([PIL.Image.new("RGB", (width, height), (128, 128, 128))] * (num_frames - 2))
-            frames.append(last_img)
-            mask_black = PIL.Image.new("L", (width, height), 0)
-            mask_white = PIL.Image.new("L", (width, height), 255)
-            mask = [mask_black, *[mask_white] * (num_frames - 2), mask_black]
-            return frames, mask
-
-        >>> # Available checkpoints: Wan-AI/Wan2.1-VACE-1.3B-diffusers, Wan-AI/Wan2.1-VACE-14B-diffusers
-        >>> model_id = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
-        >>> vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
-        >>> pipe = WanVACEPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
-        >>> flow_shift = 3.0  # 5.0 for 720P, 3.0 for 480P
-        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
-        >>> pipe.to("cuda")
-
-        >>> prompt = "CG animation style, a small blue bird takes off from the ground, flapping its wings. The bird's feathers are delicate, with a unique pattern on its chest. The background shows a blue sky with white clouds under bright sunshine. The camera follows the bird upward, capturing its flight and the vastness of the sky from a close-up, low-angle perspective."
-        >>> negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-        >>> first_frame = load_image(
-        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png"
-        ... )
-        >>> last_frame = load_image(
-        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png>>> "
-        ... )
-
-        >>> height = 512
-        >>> width = 512
-        >>> num_frames = 81
-        >>> video, mask = prepare_video_and_mask(first_frame, last_frame, height, width, num_frames)
-
-        >>> output = pipe(
-        ...     video=video,
-        ...     mask=mask,
-        ...     prompt=prompt,
-        ...     negative_prompt=negative_prompt,
-        ...     height=height,
-        ...     width=width,
-        ...     num_frames=num_frames,
-        ...     num_inference_steps=30,
-        ...     guidance_scale=5.0,
-        ...     generator=torch.Generator().manual_seed(42),
-        ... ).frames[0]
-        >>> export_to_video(output, "output.mp4", fps=16)
-        ```
-"""
-
-
-def basic_clean(text):
-    text = ftfy.fix_text(text)
-    text = html.unescape(html.unescape(text))
-    return text.strip()
-
-
-def whitespace_clean(text):
-    text = re.sub(r"\s+", " ", text)
-    text = text.strip()
-    return text
-
-
-def prompt_clean(text):
-    text = whitespace_clean(basic_clean(text))
-    return text
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
-def retrieve_latents(
-    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
-):
-    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
-        return encoder_output.latent_dist.sample(generator)
-    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
-        return encoder_output.latent_dist.mode()
-    elif hasattr(encoder_output, "latents"):
-        return encoder_output.latents
-    else:
-        raise AttributeError("Could not access latents of provided encoder_output")
-
-
-class WanVACEPipeline(DiffusionPipeline, WanLoraLoaderMixin):
-    r"""
-    Pipeline for controllable generation using Wan.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Args:
-        tokenizer ([`T5Tokenizer`]):
-            Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer),
-            specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
-        text_encoder ([`T5EncoderModel`]):
-            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
-            the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
-        transformer ([`WanTransformer3DModel`]):
-            Conditional Transformer to denoise the input latents.
-        scheduler ([`UniPCMultistepScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKLWan`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
-
-    def __init__(
-        self,
-        tokenizer: AutoTokenizer,
-        text_encoder: UMT5EncoderModel,
-        transformer: WanVACETransformer3DModel,
-        vae: AutoencoderKLWan,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-
-        self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
-        self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
-
-    # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
-    def _get_t5_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        num_videos_per_prompt: int = 1,
-        max_sequence_length: int = 226,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        prompt = [prompt_clean(u) for u in prompt]
-        batch_size = len(prompt)
-
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
-            return_attention_mask=True,
-            return_tensors="pt",
-        )
-        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
-        seq_lens = mask.gt(0).sum(dim=1).long()
-
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
-        prompt_embeds = torch.stack(
-            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
-        )
-
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
-        return prompt_embeds
-
-    # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        do_classifier_free_guidance: bool = True,
-        num_videos_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 226,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                Whether to use classifier free guidance or not.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            device: (`torch.device`, *optional*):
-                torch device
-            dtype: (`torch.dtype`, *optional*):
-                torch dtype
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        if prompt is not None:
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            negative_prompt = negative_prompt or ""
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
-
-            if prompt is not None and type(prompt) is not type(negative_prompt):
-                raise TypeError(
-                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    f" {type(prompt)}."
-                )
-            elif batch_size != len(negative_prompt):
-                raise ValueError(
-                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
-                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
-                    " the batch size of `prompt`."
-                )
-
-            negative_prompt_embeds = self._get_t5_prompt_embeds(
-                prompt=negative_prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
-                max_sequence_length=max_sequence_length,
-                device=device,
-                dtype=dtype,
-            )
-
-        return prompt_embeds, negative_prompt_embeds
-
-    def check_inputs(
-        self,
-        prompt,
-        negative_prompt,
-        height,
-        width,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        callback_on_step_end_tensor_inputs=None,
-        video=None,
-        mask=None,
-        reference_images=None,
-    ):
-        base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1]
-        if height % base != 0 or width % base != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by {base} but are {height} and {width}.")
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-        elif negative_prompt is not None and (
-            not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
-        ):
-            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
-
-        if video is not None:
-            if mask is not None:
-                if len(video) != len(mask):
-                    raise ValueError(
-                        f"Length of `video` {len(video)} and `mask` {len(mask)} do not match. Please make sure that"
-                        " they have the same length."
-                    )
-            if reference_images is not None:
-                is_pil_image = isinstance(reference_images, PIL.Image.Image)
-                is_list_of_pil_images = isinstance(reference_images, list) and all(
-                    isinstance(ref_img, PIL.Image.Image) for ref_img in reference_images
-                )
-                is_list_of_list_of_pil_images = isinstance(reference_images, list) and all(
-                    isinstance(ref_img, list) and all(isinstance(ref_img_, PIL.Image.Image) for ref_img_ in ref_img)
-                    for ref_img in reference_images
-                )
-                if not (is_pil_image or is_list_of_pil_images or is_list_of_list_of_pil_images):
-                    raise ValueError(
-                        "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, or "
-                        "`list` of `list` of `PIL.Image.Image`, but is {type(reference_images)}"
-                    )
-                if is_list_of_list_of_pil_images and len(reference_images) != 1:
-                    raise ValueError(
-                        "The pipeline only supports generating one video at a time at the moment. When passing a list "
-                        "of list of reference images, where the outer list corresponds to the batch size and the inner "
-                        "list corresponds to list of conditioning images per video, please make sure to only pass "
-                        "one inner list of reference images (i.e., `[[<image1>, <image2>, ...]]`"
-                    )
-        elif mask is not None:
-            raise ValueError("`mask` can only be passed if `video` is passed as well.")
-
-    def preprocess_conditions(
-        self,
-        video: Optional[List[PipelineImageInput]] = None,
-        mask: Optional[List[PipelineImageInput]] = None,
-        reference_images: Optional[Union[PIL.Image.Image, List[PIL.Image.Image], List[List[PIL.Image.Image]]]] = None,
-        batch_size: int = 1,
-        height: int = 480,
-        width: int = 832,
-        num_frames: int = 81,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-    ):
-        if video is not None:
-            base = self.vae_scale_factor_spatial * self.transformer.config.patch_size[1]
-            video_height, video_width = self.video_processor.get_default_height_width(video[0])
-
-            if video_height * video_width > height * width:
-                scale = min(width / video_width, height / video_height)
-                video_height, video_width = int(video_height * scale), int(video_width * scale)
-
-            if video_height % base != 0 or video_width % base != 0:
-                logger.warning(
-                    f"Video height and width should be divisible by {base}, but got {video_height} and {video_width}. "
-                )
-                video_height = (video_height // base) * base
-                video_width = (video_width // base) * base
-
-            assert video_height * video_width <= height * width
-
-            video = self.video_processor.preprocess_video(video, video_height, video_width)
-            image_size = (video_height, video_width)  # Use the height/width of video (with possible rescaling)
-        else:
-            video = torch.zeros(batch_size, 3, num_frames, height, width, dtype=dtype, device=device)
-            image_size = (height, width)  # Use the height/width provider by user
-
-        if mask is not None:
-            mask = self.video_processor.preprocess_video(mask, image_size[0], image_size[1])
-            mask = torch.clamp((mask + 1) / 2, min=0, max=1)
-        else:
-            mask = torch.ones_like(video)
-
-        video = video.to(dtype=dtype, device=device)
-        mask = mask.to(dtype=dtype, device=device)
-
-        # Make a list of list of images where the outer list corresponds to video batch size and the inner list
-        # corresponds to list of conditioning images per video
-        if reference_images is None or isinstance(reference_images, PIL.Image.Image):
-            reference_images = [[reference_images] for _ in range(video.shape[0])]
-        elif isinstance(reference_images, (list, tuple)) and isinstance(next(iter(reference_images)), PIL.Image.Image):
-            reference_images = [reference_images]
-        elif (
-            isinstance(reference_images, (list, tuple))
-            and isinstance(next(iter(reference_images)), list)
-            and isinstance(next(iter(reference_images[0])), PIL.Image.Image)
-        ):
-            reference_images = reference_images
-        else:
-            raise ValueError(
-                "`reference_images` has to be of type `PIL.Image.Image` or `list` of `PIL.Image.Image`, or "
-                "`list` of `list` of `PIL.Image.Image`, but is {type(reference_images)}"
-            )
-
-        if video.shape[0] != len(reference_images):
-            raise ValueError(
-                f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match."
-            )
-
-        ref_images_lengths = [len(reference_images_batch) for reference_images_batch in reference_images]
-        if any(l != ref_images_lengths[0] for l in ref_images_lengths):
-            raise ValueError(
-                f"All batches of `reference_images` should have the same length, but got {ref_images_lengths}. Support for this "
-                "may be added in the future."
-            )
-
-        reference_images_preprocessed = []
-        for i, reference_images_batch in enumerate(reference_images):
-            preprocessed_images = []
-            for j, image in enumerate(reference_images_batch):
-                if image is None:
-                    continue
-                image = self.video_processor.preprocess(image, None, None)
-                img_height, img_width = image.shape[-2:]
-                scale = min(image_size[0] / img_height, image_size[1] / img_width)
-                new_height, new_width = int(img_height * scale), int(img_width * scale)
-                resized_image = torch.nn.functional.interpolate(
-                    image, size=(new_height, new_width), mode="bilinear", align_corners=False
-                ).squeeze(0)  # [C, H, W]
-                top = (image_size[0] - new_height) // 2
-                left = (image_size[1] - new_width) // 2
-                canvas = torch.ones(3, *image_size, device=device, dtype=dtype)
-                canvas[:, top : top + new_height, left : left + new_width] = resized_image
-                preprocessed_images.append(canvas)
-            reference_images_preprocessed.append(preprocessed_images)
-
-        return video, mask, reference_images_preprocessed
-
-    def prepare_video_latents(
-        self,
-        video: torch.Tensor,
-        mask: torch.Tensor,
-        reference_images: Optional[List[List[torch.Tensor]]] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        device: Optional[torch.device] = None,
-    ) -> torch.Tensor:
-        device = device or self._execution_device
-
-        if isinstance(generator, list):
-            # TODO: support this
-            raise ValueError("Passing a list of generators is not yet supported. This may be supported in the future.")
-
-        if reference_images is None:
-            # For each batch of video, we set no re
-            # ference image (as one or more can be passed by user)
-            reference_images = [[None] for _ in range(video.shape[0])]
-        else:
-            if video.shape[0] != len(reference_images):
-                raise ValueError(
-                    f"Batch size of `video` {video.shape[0]} and length of `reference_images` {len(reference_images)} does not match."
-                )
-
-        if video.shape[0] != 1:
-            # TODO: support this
-            raise ValueError(
-                "Generating with more than one video is not yet supported. This may be supported in the future."
-            )
-
-        vae_dtype = self.vae.dtype
-        video = video.to(dtype=vae_dtype)
-
-        latents_mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=torch.float32).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        )
-        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std, device=device, dtype=torch.float32).view(
-            1, self.vae.config.z_dim, 1, 1, 1
-        )
-
-        if mask is None:
-            latents = retrieve_latents(self.vae.encode(video), generator, sample_mode="argmax").unbind(0)
-            latents = ((latents.float() - latents_mean) * latents_std).to(vae_dtype)
-        else:
-            mask = mask.to(dtype=vae_dtype)
-            mask = torch.where(mask > 0.5, 1.0, 0.0)
-            inactive = video * (1 - mask)
-            reactive = video * mask
-            inactive = retrieve_latents(self.vae.encode(inactive), generator, sample_mode="argmax")
-            reactive = retrieve_latents(self.vae.encode(reactive), generator, sample_mode="argmax")
-            inactive = ((inactive.float() - latents_mean) * latents_std).to(vae_dtype)
-            reactive = ((reactive.float() - latents_mean) * latents_std).to(vae_dtype)
-            latents = torch.cat([inactive, reactive], dim=1)
-
-        latent_list = []
-        for latent, reference_images_batch in zip(latents, reference_images):
-            for reference_image in reference_images_batch:
-                assert reference_image.ndim == 3
-                reference_image = reference_image.to(dtype=vae_dtype)
-                reference_image = reference_image[None, :, None, :, :]  # [1, C, 1, H, W]
-                reference_latent = retrieve_latents(self.vae.encode(reference_image), generator, sample_mode="argmax")
-                reference_latent = ((reference_latent.float() - latents_mean) * latents_std).to(vae_dtype)
-                reference_latent = reference_latent.squeeze(0)  # [C, 1, H, W]
-                reference_latent = torch.cat([reference_latent, torch.zeros_like(reference_latent)], dim=0)
-                latent = torch.cat([reference_latent.squeeze(0), latent], dim=1)
-            latent_list.append(latent)
-        return torch.stack(latent_list)
-
-    def prepare_masks(
-        self,
-        mask: torch.Tensor,
-        reference_images: Optional[List[torch.Tensor]] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-    ) -> torch.Tensor:
-        if isinstance(generator, list):
-            # TODO: support this
-            raise ValueError("Passing a list of generators is not yet supported. This may be supported in the future.")
-
-        if reference_images is None:
-            # For each batch of video, we set no reference image (as one or more can be passed by user)
-            reference_images = [[None] for _ in range(mask.shape[0])]
-        else:
-            if mask.shape[0] != len(reference_images):
-                raise ValueError(
-                    f"Batch size of `mask` {mask.shape[0]} and length of `reference_images` {len(reference_images)} does not match."
-                )
-
-        if mask.shape[0] != 1:
-            # TODO: support this
-            raise ValueError(
-                "Generating with more than one video is not yet supported. This may be supported in the future."
-            )
-
-        transformer_patch_size = self.transformer.config.patch_size[1]
-
-        mask_list = []
-        for mask_, reference_images_batch in zip(mask, reference_images):
-            num_channels, num_frames, height, width = mask_.shape
-            new_num_frames = (num_frames + self.vae_scale_factor_temporal - 1) // self.vae_scale_factor_temporal
-            new_height = height // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size
-            new_width = width // (self.vae_scale_factor_spatial * transformer_patch_size) * transformer_patch_size
-            mask_ = mask_[0, :, :, :]
-            mask_ = mask_.view(
-                num_frames, new_height, self.vae_scale_factor_spatial, new_width, self.vae_scale_factor_spatial
-            )
-            mask_ = mask_.permute(2, 4, 0, 1, 3).flatten(0, 1)  # [8x8, num_frames, new_height, new_width]
-            mask_ = torch.nn.functional.interpolate(
-                mask_.unsqueeze(0), size=(new_num_frames, new_height, new_width), mode="nearest-exact"
-            ).squeeze(0)
-            num_ref_images = len(reference_images_batch)
-            if num_ref_images > 0:
-                mask_padding = torch.zeros_like(mask_[:, :num_ref_images, :, :])
-                mask_ = torch.cat([mask_, mask_padding], dim=1)
-            mask_list.append(mask_)
-        return torch.stack(mask_list)
-
-    def prepare_latents(
-        self,
-        batch_size: int,
-        num_channels_latents: int = 16,
-        height: int = 480,
-        width: int = 832,
-        num_frames: int = 81,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            return latents.to(device=device, dtype=dtype)
-
-        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        shape = (
-            batch_size,
-            num_channels_latents,
-            num_latent_frames,
-            int(height) // self.vae_scale_factor_spatial,
-            int(width) // self.vae_scale_factor_spatial,
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        return latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1.0
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def current_timestep(self):
-        return self._current_timestep
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Union[str, List[str]] = None,
-        video: Optional[List[PipelineImageInput]] = None,
-        mask: Optional[List[PipelineImageInput]] = None,
-        reference_images: Optional[List[PipelineImageInput]] = None,
-        conditioning_scale: Union[float, List[float], torch.Tensor] = 1.0,
-        height: int = 480,
-        width: int = 832,
-        num_frames: int = 81,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 5.0,
-        num_videos_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "np",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[
-            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
-        ] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ):
-        r"""
-        The call function to the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            height (`int`, defaults to `480`):
-                The height in pixels of the generated image.
-            width (`int`, defaults to `832`):
-                The width in pixels of the generated image.
-            num_frames (`int`, defaults to `81`):
-                The number of frames in the generated video.
-            num_inference_steps (`int`, defaults to `50`):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            guidance_scale (`float`, defaults to `5.0`):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            num_videos_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
-                generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor is generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
-                provided, text embeddings are generated from the `prompt` input argument.
-            output_type (`str`, *optional*, defaults to `"np"`):
-                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`WanPipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
-                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
-                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
-                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
-                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            autocast_dtype (`torch.dtype`, *optional*, defaults to `torch.bfloat16`):
-                The dtype to use for the torch.amp.autocast.
-
-        Examples:
-
-        Returns:
-            [`~WanPipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`WanPipelineOutput`] is returned, otherwise a `tuple` is returned where
-                the first element is a list with the generated images and the second element is a list of `bool`s
-                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
-        """
-
-        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
-            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
-
-        # Simplification of implementation for now
-        if not isinstance(prompt, str):
-            raise ValueError("Passing a list of prompts is not yet supported. This may be supported in the future.")
-        if num_videos_per_prompt != 1:
-            raise ValueError(
-                "Generating multiple videos per prompt is not yet supported. This may be supported in the future."
-            )
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            negative_prompt,
-            height,
-            width,
-            prompt_embeds,
-            negative_prompt_embeds,
-            callback_on_step_end_tensor_inputs,
-            video,
-            mask,
-            reference_images,
-        )
-
-        if num_frames % self.vae_scale_factor_temporal != 1:
-            logger.warning(
-                f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number."
-            )
-            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
-        num_frames = max(num_frames, 1)
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._current_timestep = None
-        self._interrupt = False
-
-        device = self._execution_device
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        vae_dtype = self.vae.dtype
-        transformer_dtype = self.transformer.dtype
-
-        if isinstance(conditioning_scale, (int, float)):
-            conditioning_scale = [conditioning_scale] * len(self.transformer.config.vace_layers)
-        if isinstance(conditioning_scale, list):
-            if len(conditioning_scale) != len(self.transformer.config.vace_layers):
-                raise ValueError(
-                    f"Length of `conditioning_scale` {len(conditioning_scale)} does not match number of layers {len(self.transformer.config.vace_layers)}."
-                )
-            conditioning_scale = torch.tensor(conditioning_scale)
-        if isinstance(conditioning_scale, torch.Tensor):
-            if conditioning_scale.size(0) != len(self.transformer.config.vace_layers):
-                raise ValueError(
-                    f"Length of `conditioning_scale` {conditioning_scale.size(0)} does not match number of layers {len(self.transformer.config.vace_layers)}."
-                )
-            conditioning_scale = conditioning_scale.to(device=device, dtype=transformer_dtype)
-
-        # 3. Encode input prompt
-        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
-            num_videos_per_prompt=num_videos_per_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            max_sequence_length=max_sequence_length,
-            device=device,
-        )
-
-        prompt_embeds = prompt_embeds.to(transformer_dtype)
-        if negative_prompt_embeds is not None:
-            negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = self.scheduler.timesteps
-
-        # 5. Prepare latent variables
-        video, mask, reference_images = self.preprocess_conditions(
-            video,
-            mask,
-            reference_images,
-            batch_size,
-            height,
-            width,
-            num_frames,
-            torch.float32,
-            device,
-        )
-        num_reference_images = len(reference_images[0])
-
-        conditioning_latents = self.prepare_video_latents(video, mask, reference_images, generator, device)
-        mask = self.prepare_masks(mask, reference_images, generator)
-        conditioning_latents = torch.cat([conditioning_latents, mask], dim=1)
-        conditioning_latents = conditioning_latents.to(transformer_dtype)
-
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames + num_reference_images * self.vae_scale_factor_temporal,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-
-        if conditioning_latents.shape[2] != latents.shape[2]:
-            logger.warning(
-                "The number of frames in the conditioning latents does not match the number of frames to be generated. Generation quality may be affected."
-            )
-
-        # 6. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-        self._num_timesteps = len(timesteps)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                self._current_timestep = t
-                latent_model_input = latents.to(transformer_dtype)
-                timestep = t.expand(latents.shape[0])
-
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    encoder_hidden_states=prompt_embeds,
-                    control_hidden_states=conditioning_latents,
-                    control_hidden_states_scale=conditioning_scale,
-                    attention_kwargs=attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                if self.do_classifier_free_guidance:
-                    noise_uncond = self.transformer(
-                        hidden_states=latent_model_input,
-                        timestep=timestep,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        control_hidden_states=conditioning_latents,
-                        control_hidden_states_scale=conditioning_scale,
-                        attention_kwargs=attention_kwargs,
-                        return_dict=False,
-                    )[0]
-                    noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        self._current_timestep = None
-
-        if not output_type == "latent":
-            latents = latents[:, :, num_reference_images:]
-            latents = latents.to(vae_dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-        else:
-            video = latents
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (video,)
-
-        return WanPipelineOutput(frames=video)

src/diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py

@@ -87,7 +87,6 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         lower_order_final: bool = True,
         euler_at_final: bool = False,
         final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
-        use_flow_sigmas: bool = False,
     ):
         if solver_type not in ["midpoint", "heun"]:
             if solver_type in ["logrho", "bh1", "bh2"]:
@@ -153,19 +152,23 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         if not isinstance(sigma, torch.Tensor):
             sigma = torch.tensor([sigma])
 
-        if self.config.use_flow_sigmas:
-            c_noise = sigma / (sigma + 1)
-        else:
-            c_noise = sigma.atan() / math.pi * 2
-
-        return c_noise
+        return sigma.atan() / math.pi * 2
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs
     def precondition_outputs(self, sample, model_output, sigma):
-        if self.config.use_flow_sigmas:
-            return self._precondition_outputs_flow(sample, model_output, sigma)
+        sigma_data = self.config.sigma_data
+        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
+
+        if self.config.prediction_type == "epsilon":
+            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
+        elif self.config.prediction_type == "v_prediction":
+            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
         else:
-            return self._precondition_outputs_edm(sample, model_output, sigma)
+            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
+
+        denoised = c_skip * sample + c_out * model_output
+
+        return denoised
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input
     def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
@@ -567,42 +570,8 @@ class CosineDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in
     def _get_conditioning_c_in(self, sigma):
-        if self.config.use_flow_sigmas:
-            t = sigma / (sigma + 1)
-            c_in = 1.0 - t
-        else:
-            c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
+        c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
         return c_in
 
-    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._precondition_outputs_flow
-    def _precondition_outputs_flow(self, sample, model_output, sigma):
-        t = sigma / (sigma + 1)
-        c_skip = 1.0 - t
-
-        if self.config.prediction_type == "epsilon":
-            c_out = -t
-        elif self.config.prediction_type == "v_prediction":
-            c_out = t
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
-    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._precondition_outputs_edm
-    def _precondition_outputs_edm(self, sample, model_output, sigma):
-        sigma_data = self.config.sigma_data
-        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
-
-        if self.config.prediction_type == "epsilon":
-            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        elif self.config.prediction_type == "v_prediction":
-            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
     def __len__(self):
         return self.config.num_train_timesteps

src/diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py

@@ -15,35 +15,14 @@
 # DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver and https://github.com/NVlabs/edm
 
 import math
-from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
 
 from ..configuration_utils import ConfigMixin, register_to_config
-from ..utils import BaseOutput
 from ..utils.torch_utils import randn_tensor
-from .scheduling_utils import SchedulerMixin
-
-
-@dataclass
-# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EDMDPMSolverMultistep
-class EDMDPMSolverMultistepSchedulerOutput(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-
-    Args:
-        prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-        pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images):
-            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
-            `pred_original_sample` can be used to preview progress or for guidance.
-    """
-
-    prev_sample: torch.Tensor
-    pred_original_sample: Optional[torch.Tensor] = None
+from .scheduling_utils import SchedulerMixin, SchedulerOutput
 
 
 class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
@@ -128,7 +107,6 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         lower_order_final: bool = True,
         euler_at_final: bool = False,
         final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
-        use_flow_sigmas: bool = False,
     ):
         # settings for DPM-Solver
         if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"]:
@@ -207,19 +185,25 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         if not isinstance(sigma, torch.Tensor):
             sigma = torch.tensor([sigma])
 
-        if self.config.use_flow_sigmas:
-            c_noise = sigma / (sigma + 1)
-        else:
-            c_noise = 0.25 * torch.log(sigma)
+        c_noise = 0.25 * torch.log(sigma)
 
         return c_noise
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.precondition_outputs
     def precondition_outputs(self, sample, model_output, sigma):
-        if self.config.use_flow_sigmas:
-            return self._precondition_outputs_flow(sample, model_output, sigma)
+        sigma_data = self.config.sigma_data
+        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
+
+        if self.config.prediction_type == "epsilon":
+            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
+        elif self.config.prediction_type == "v_prediction":
+            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
         else:
-            return self._precondition_outputs_edm(sample, model_output, sigma)
+            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
+
+        denoised = c_skip * sample + c_out * model_output
+
+        return denoised
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler.scale_model_input
     def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
@@ -614,8 +598,7 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         sample: torch.Tensor,
         generator=None,
         return_dict: bool = True,
-        pred_original_sample: Optional[torch.Tensor] = None,
-    ) -> Union[EDMDPMSolverMultistepSchedulerOutput, Tuple]:
+    ) -> Union[SchedulerOutput, Tuple]:
         """
         Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
         the multistep DPMSolver.
@@ -630,14 +613,12 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
             generator (`torch.Generator`, *optional*):
                 A random number generator.
             return_dict (`bool`):
-                Whether or not to return a
-                [`~schedulers.scheduling_edm_dpmsolver_multistep.EDMDPMSolverMultistepSchedulerOutput`] or a `tuple`.
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
 
         Returns:
-            [`~schedulers.scheduling_edm_dpmsolver_multistep.EDMDPMSolverMultistepSchedulerOutput`] or `tuple`:
-                If return_dict is `True`,
-                [`~schedulers.scheduling_edm_dpmsolver_multistep.EDMDPMSolverMultistepSchedulerOutput`] is returned,
-                otherwise a tuple is returned where the first element is the sample tensor.
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+                tuple is returned where the first element is the sample tensor.
 
         """
         if self.num_inference_steps is None:
@@ -658,12 +639,7 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
             (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
         )
 
-        if pred_original_sample is None:
-            model_output = self.convert_model_output(model_output, sample=sample)
-        else:
-            model_output = pred_original_sample
-            # TODO: thresholding is not handled in this case, but probably not needed either for Cosmos
-
+        model_output = self.convert_model_output(model_output, sample=sample)
         for i in range(self.config.solver_order - 1):
             self.model_outputs[i] = self.model_outputs[i + 1]
         self.model_outputs[-1] = model_output
@@ -691,7 +667,7 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
         if not return_dict:
             return (prev_sample,)
 
-        return EDMDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample, pred_original_sample=model_output)
+        return SchedulerOutput(prev_sample=prev_sample)
 
     # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler.add_noise
     def add_noise(
@@ -729,42 +705,8 @@ class EDMDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
 
     # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._get_conditioning_c_in
     def _get_conditioning_c_in(self, sigma):
-        if self.config.use_flow_sigmas:
-            t = sigma / (sigma + 1)
-            c_in = 1.0 - t
-        else:
-            c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
+        c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
         return c_in
 
-    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._precondition_outputs_flow
-    def _precondition_outputs_flow(self, sample, model_output, sigma):
-        t = sigma / (sigma + 1)
-        c_skip = 1.0 - t
-
-        if self.config.prediction_type == "epsilon":
-            c_out = -t
-        elif self.config.prediction_type == "v_prediction":
-            c_out = t
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
-    # Copied from diffusers.schedulers.scheduling_edm_euler.EDMEulerScheduler._precondition_outputs_edm
-    def _precondition_outputs_edm(self, sample, model_output, sigma):
-        sigma_data = self.config.sigma_data
-        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
-
-        if self.config.prediction_type == "epsilon":
-            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        elif self.config.prediction_type == "v_prediction":
-            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
     def __len__(self):
         return self.config.num_train_timesteps

src/diffusers/schedulers/scheduling_edm_euler.py

@@ -28,7 +28,7 @@ logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
 @dataclass
-# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EDMEuler
+# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->EulerDiscrete
 class EDMEulerSchedulerOutput(BaseOutput):
     """
     Output class for the scheduler's `step` function output.
@@ -96,7 +96,6 @@ class EDMEulerScheduler(SchedulerMixin, ConfigMixin):
         prediction_type: str = "epsilon",
         rho: float = 7.0,
         final_sigmas_type: str = "zero",  # can be "zero" or "sigma_min"
-        use_flow_sigmas: bool = False,
     ):
         if sigma_schedule not in ["karras", "exponential"]:
             raise ValueError(f"Wrong value for provided for `{sigma_schedule=}`.`")
@@ -170,18 +169,24 @@ class EDMEulerScheduler(SchedulerMixin, ConfigMixin):
         if not isinstance(sigma, torch.Tensor):
             sigma = torch.tensor([sigma])
 
-        if self.config.use_flow_sigmas:
-            c_noise = sigma / (sigma + 1)
-        else:
-            c_noise = 0.25 * torch.log(sigma)
+        c_noise = 0.25 * torch.log(sigma)
 
         return c_noise
 
     def precondition_outputs(self, sample, model_output, sigma):
-        if self.config.use_flow_sigmas:
-            return self._precondition_outputs_flow(sample, model_output, sigma)
+        sigma_data = self.config.sigma_data
+        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
+
+        if self.config.prediction_type == "epsilon":
+            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
+        elif self.config.prediction_type == "v_prediction":
+            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
         else:
-            return self._precondition_outputs_edm(sample, model_output, sigma)
+            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
+
+        denoised = c_skip * sample + c_out * model_output
+
+        return denoised
 
     def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
         """
@@ -436,40 +441,8 @@ class EDMEulerScheduler(SchedulerMixin, ConfigMixin):
         return noisy_samples
 
     def _get_conditioning_c_in(self, sigma):
-        if self.config.use_flow_sigmas:
-            t = sigma / (sigma + 1)
-            c_in = 1.0 - t
-        else:
-            c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
+        c_in = 1 / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
         return c_in
 
-    def _precondition_outputs_flow(self, sample, model_output, sigma):
-        t = sigma / (sigma + 1)
-        c_skip = 1.0 - t
-
-        if self.config.prediction_type == "epsilon":
-            c_out = -t
-        elif self.config.prediction_type == "v_prediction":
-            c_out = t
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
-    def _precondition_outputs_edm(self, sample, model_output, sigma):
-        sigma_data = self.config.sigma_data
-        c_skip = sigma_data**2 / (sigma**2 + sigma_data**2)
-
-        if self.config.prediction_type == "epsilon":
-            c_out = sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        elif self.config.prediction_type == "v_prediction":
-            c_out = -sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
-        else:
-            raise ValueError(f"Prediction type {self.config.prediction_type} is not supported.")
-
-        denoised = c_skip * sample + c_out * model_output
-        return denoised
-
     def __len__(self):
         return self.config.num_train_timesteps

src/diffusers/utils/dummy_pt_objects.py

@@ -1150,21 +1150,6 @@ class WanTransformer3DModel(metaclass=DummyObject):
         requires_backends(cls, ["torch"])
 
 
-class WanVACETransformer3DModel(metaclass=DummyObject):
-    _backends = ["torch"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch"])
-
-
 def get_constant_schedule(*args, **kwargs):
     requires_backends(get_constant_schedule, ["torch"])
 

src/diffusers/utils/dummy_torch_and_transformers_objects.py

@@ -407,36 +407,6 @@ class ConsisIDPipeline(metaclass=DummyObject):
         requires_backends(cls, ["torch", "transformers"])
 
 
-class Cosmos2TextToImagePipeline(metaclass=DummyObject):
-    _backends = ["torch", "transformers"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch", "transformers"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-
-class Cosmos2VideoToWorldPipeline(metaclass=DummyObject):
-    _backends = ["torch", "transformers"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch", "transformers"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-
 class CosmosTextToWorldPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 
@@ -1652,21 +1622,6 @@ class SanaPipeline(metaclass=DummyObject):
         requires_backends(cls, ["torch", "transformers"])
 
 
-class SanaSprintImg2ImgPipeline(metaclass=DummyObject):
-    _backends = ["torch", "transformers"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch", "transformers"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-
 class SanaSprintPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 
@@ -2927,21 +2882,6 @@ class WanPipeline(metaclass=DummyObject):
         requires_backends(cls, ["torch", "transformers"])
 
 
-class WanVACEPipeline(metaclass=DummyObject):
-    _backends = ["torch", "transformers"]
-
-    def __init__(self, *args, **kwargs):
-        requires_backends(self, ["torch", "transformers"])
-
-    @classmethod
-    def from_config(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        requires_backends(cls, ["torch", "transformers"])
-
-
 class WanVideoToVideoPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 

src/diffusers/utils/torch_utils.py

@@ -18,7 +18,7 @@ PyTorch utilities: Utilities related to PyTorch
 from typing import List, Optional, Tuple, Union
 
 from . import logging
-from .import_utils import is_torch_available, is_torch_npu_available, is_torch_version
+from .import_utils import is_torch_available, is_torch_version
 
 
 if is_torch_available():
@@ -166,8 +166,6 @@ def get_torch_cuda_device_capability():
 def get_device():
     if torch.cuda.is_available():
         return "cuda"
-    elif is_torch_npu_available():
-        return "npu"
     elif hasattr(torch, "xpu") and torch.xpu.is_available():
         return "xpu"
     else:

tests/hooks/test_group_offloading.py

@@ -22,13 +22,7 @@ from diffusers.models import ModelMixin
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.utils import get_logger
 from diffusers.utils.import_utils import compare_versions
-from diffusers.utils.testing_utils import (
-    backend_empty_cache,
-    backend_max_memory_allocated,
-    backend_reset_peak_memory_stats,
-    require_torch_accelerator,
-    torch_device,
-)
+from diffusers.utils.testing_utils import require_torch_gpu, torch_device
 
 
 class DummyBlock(torch.nn.Module):
@@ -113,7 +107,7 @@ class DummyPipeline(DiffusionPipeline):
         return x
 
 
-@require_torch_accelerator
+@require_torch_gpu
 class GroupOffloadTests(unittest.TestCase):
     in_features = 64
     hidden_features = 256
@@ -131,8 +125,8 @@ class GroupOffloadTests(unittest.TestCase):
         del self.model
         del self.input
         gc.collect()
-        backend_empty_cache(torch_device)
-        backend_reset_peak_memory_stats(torch_device)
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
 
     def get_model(self):
         torch.manual_seed(0)
@@ -147,8 +141,8 @@ class GroupOffloadTests(unittest.TestCase):
         @torch.no_grad()
         def run_forward(model):
             gc.collect()
-            backend_empty_cache(torch_device)
-            backend_reset_peak_memory_stats(torch_device)
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
             self.assertTrue(
                 all(
                     module._diffusers_hook.get_hook("group_offloading") is not None
@@ -158,7 +152,7 @@ class GroupOffloadTests(unittest.TestCase):
             )
             model.eval()
             output = model(self.input)[0].cpu()
-            max_memory_allocated = backend_max_memory_allocated(torch_device)
+            max_memory_allocated = torch.cuda.max_memory_allocated()
             return output, max_memory_allocated
 
         self.model.to(torch_device)
@@ -193,10 +187,10 @@ class GroupOffloadTests(unittest.TestCase):
         self.assertTrue(torch.allclose(output_without_group_offloading, output_with_group_offloading5, atol=1e-5))
 
         # Memory assertions - offloading should reduce memory usage
-        self.assertTrue(mem4 <= mem5 < mem2 <= mem3 < mem1 < mem_baseline)
+        self.assertTrue(mem4 <= mem5 < mem2 < mem3 < mem1 < mem_baseline)
 
-    def test_warning_logged_if_group_offloaded_module_moved_to_accelerator(self):
-        if torch.device(torch_device).type not in ["cuda", "xpu"]:
+    def test_warning_logged_if_group_offloaded_module_moved_to_cuda(self):
+        if torch.device(torch_device).type != "cuda":
             return
         self.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
         logger = get_logger("diffusers.models.modeling_utils")
@@ -205,8 +199,8 @@ class GroupOffloadTests(unittest.TestCase):
             self.model.to(torch_device)
         self.assertIn(f"The module '{self.model.__class__.__name__}' is group offloaded", cm.output[0])
 
-    def test_warning_logged_if_group_offloaded_pipe_moved_to_accelerator(self):
-        if torch.device(torch_device).type not in ["cuda", "xpu"]:
+    def test_warning_logged_if_group_offloaded_pipe_moved_to_cuda(self):
+        if torch.device(torch_device).type != "cuda":
             return
         pipe = DummyPipeline(self.model)
         self.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
@@ -216,20 +210,19 @@ class GroupOffloadTests(unittest.TestCase):
             pipe.to(torch_device)
         self.assertIn(f"The module '{self.model.__class__.__name__}' is group offloaded", cm.output[0])
 
-    def test_error_raised_if_streams_used_and_no_accelerator_device(self):
-        torch_accelerator_module = getattr(torch, torch_device, torch.cuda)
-        original_is_available = torch_accelerator_module.is_available
-        torch_accelerator_module.is_available = lambda: False
+    def test_error_raised_if_streams_used_and_no_cuda_device(self):
+        original_is_available = torch.cuda.is_available
+        torch.cuda.is_available = lambda: False
         with self.assertRaises(ValueError):
             self.model.enable_group_offload(
-                onload_device=torch.device(torch_device), offload_type="leaf_level", use_stream=True
+                onload_device=torch.device("cuda"), offload_type="leaf_level", use_stream=True
             )
-        torch_accelerator_module.is_available = original_is_available
+        torch.cuda.is_available = original_is_available
 
     def test_error_raised_if_supports_group_offloading_false(self):
         self.model._supports_group_offloading = False
         with self.assertRaisesRegex(ValueError, "does not support group offloading"):
-            self.model.enable_group_offload(onload_device=torch.device(torch_device))
+            self.model.enable_group_offload(onload_device=torch.device("cuda"))
 
     def test_error_raised_if_model_offloading_applied_on_group_offloaded_module(self):
         pipe = DummyPipeline(self.model)
@@ -256,7 +249,7 @@ class GroupOffloadTests(unittest.TestCase):
             pipe.model.enable_group_offload(torch_device, offload_type="block_level", num_blocks_per_group=3)
 
     def test_block_level_stream_with_invocation_order_different_from_initialization_order(self):
-        if torch.device(torch_device).type not in ["cuda", "xpu"]:
+        if torch.device(torch_device).type != "cuda":
             return
         model = DummyModelWithMultipleBlocks(
             in_features=self.in_features,

tests/lora/test_lora_layers_cogvideox.py

@@ -124,9 +124,6 @@ class CogVideoXLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests):
     def test_simple_inference_with_text_denoiser_lora_unfused(self):
         super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3)
 
-    def test_lora_scale_kwargs_match_fusion(self):
-        super().test_lora_scale_kwargs_match_fusion(expected_atol=9e-3, expected_rtol=9e-3)
-
     @unittest.skip("Not supported in CogVideoX.")
     def test_simple_inference_with_text_denoiser_block_scale(self):
         pass

tests/lora/test_lora_layers_sdxl.py

@@ -117,40 +117,6 @@ class StableDiffusionXLLoRATests(PeftLoraLoaderMixinTests, unittest.TestCase):
     def test_multiple_wrong_adapter_name_raises_error(self):
         super().test_multiple_wrong_adapter_name_raises_error()
 
-    def test_simple_inference_with_text_denoiser_lora_unfused(self):
-        if torch.cuda.is_available():
-            expected_atol = 9e-2
-            expected_rtol = 9e-2
-        else:
-            expected_atol = 1e-3
-            expected_rtol = 1e-3
-
-        super().test_simple_inference_with_text_denoiser_lora_unfused(
-            expected_atol=expected_atol, expected_rtol=expected_rtol
-        )
-
-    def test_simple_inference_with_text_lora_denoiser_fused_multi(self):
-        if torch.cuda.is_available():
-            expected_atol = 9e-2
-            expected_rtol = 9e-2
-        else:
-            expected_atol = 1e-3
-            expected_rtol = 1e-3
-
-        super().test_simple_inference_with_text_lora_denoiser_fused_multi(
-            expected_atol=expected_atol, expected_rtol=expected_rtol
-        )
-
-    def test_lora_scale_kwargs_match_fusion(self):
-        if torch.cuda.is_available():
-            expected_atol = 9e-2
-            expected_rtol = 9e-2
-        else:
-            expected_atol = 1e-3
-            expected_rtol = 1e-3
-
-        super().test_lora_scale_kwargs_match_fusion(expected_atol=expected_atol, expected_rtol=expected_rtol)
-
 
 @slow
 @nightly

tests/lora/utils.py

@@ -80,18 +80,6 @@ def initialize_dummy_state_dict(state_dict):
 POSSIBLE_ATTENTION_KWARGS_NAMES = ["cross_attention_kwargs", "joint_attention_kwargs", "attention_kwargs"]
 
 
-def determine_attention_kwargs_name(pipeline_class):
-    call_signature_keys = inspect.signature(pipeline_class.__call__).parameters.keys()
-
-    # TODO(diffusers): Discuss a common naming convention across library for 1.0.0 release
-    for possible_attention_kwargs in POSSIBLE_ATTENTION_KWARGS_NAMES:
-        if possible_attention_kwargs in call_signature_keys:
-            attention_kwargs_name = possible_attention_kwargs
-            break
-    assert attention_kwargs_name is not None
-    return attention_kwargs_name
-
-
 @require_peft_backend
 class PeftLoraLoaderMixinTests:
     pipeline_class = None
@@ -454,7 +442,14 @@ class PeftLoraLoaderMixinTests:
         Tests a simple inference with lora attached on the text encoder + scale argument
         and makes sure it works as expected
         """
-        attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class)
+        call_signature_keys = inspect.signature(self.pipeline_class.__call__).parameters.keys()
+
+        # TODO(diffusers): Discuss a common naming convention across library for 1.0.0 release
+        for possible_attention_kwargs in POSSIBLE_ATTENTION_KWARGS_NAMES:
+            if possible_attention_kwargs in call_signature_keys:
+                attention_kwargs_name = possible_attention_kwargs
+                break
+        assert attention_kwargs_name is not None
 
         for scheduler_cls in self.scheduler_classes:
             components, text_lora_config, _ = self.get_dummy_components(scheduler_cls)
@@ -745,7 +740,12 @@ class PeftLoraLoaderMixinTests:
         Tests a simple inference with lora attached on the text encoder + Unet + scale argument
         and makes sure it works as expected
         """
-        attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class)
+        call_signature_keys = inspect.signature(self.pipeline_class.__call__).parameters.keys()
+        for possible_attention_kwargs in POSSIBLE_ATTENTION_KWARGS_NAMES:
+            if possible_attention_kwargs in call_signature_keys:
+                attention_kwargs_name = possible_attention_kwargs
+                break
+        assert attention_kwargs_name is not None
 
         for scheduler_cls in self.scheduler_classes:
             components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
@@ -878,11 +878,9 @@ class PeftLoraLoaderMixinTests:
             pipe, denoiser = self.check_if_adapters_added_correctly(pipe, text_lora_config, denoiser_lora_config)
 
             pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules)
-            self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
             output_fused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
             pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules)
-            self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
             output_unfused_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
             # unloading should remove the LoRA layers
@@ -1092,7 +1090,7 @@ class PeftLoraLoaderMixinTests:
     def test_simple_inference_with_text_denoiser_multi_adapter_block_lora(self):
         """
         Tests a simple inference with lora attached to text encoder and unet, attaches
-        multiple adapters and set different weights for different blocks (i.e. block lora)
+        multiple adapters and set differnt weights for different blocks (i.e. block lora)
         """
         for scheduler_cls in self.scheduler_classes:
             components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
@@ -1610,21 +1608,26 @@ class PeftLoraLoaderMixinTests:
                 self.assertTrue(
                     check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder"
                 )
-                pipe.text_encoder.add_adapter(text_lora_config, "adapter-2")
 
             denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
             denoiser.add_adapter(denoiser_lora_config, "adapter-1")
-            self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
+
+            # Attach a second adapter
+            if "text_encoder" in self.pipeline_class._lora_loadable_modules:
+                pipe.text_encoder.add_adapter(text_lora_config, "adapter-2")
+
             denoiser.add_adapter(denoiser_lora_config, "adapter-2")
 
+            self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
+
             if self.has_two_text_encoders or self.has_three_text_encoders:
                 lora_loadable_components = self.pipeline_class._lora_loadable_modules
                 if "text_encoder_2" in lora_loadable_components:
                     pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1")
+                    pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2")
                     self.assertTrue(
                         check_if_lora_correctly_set(pipe.text_encoder_2), "Lora not correctly set in text encoder 2"
                     )
-                    pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-2")
 
             # set them to multi-adapter inference mode
             pipe.set_adapters(["adapter-1", "adapter-2"])
@@ -1634,9 +1637,8 @@ class PeftLoraLoaderMixinTests:
             outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
             pipe.fuse_lora(components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-1"])
-            self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
 
-            # Fusing should still keep the LoRA layers so output should remain the same
+            # Fusing should still keep the LoRA layers so outpout should remain the same
             outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0]
 
             self.assertTrue(
@@ -1645,23 +1647,9 @@ class PeftLoraLoaderMixinTests:
             )
 
             pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules)
-            self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
-
-            if "text_encoder" in self.pipeline_class._lora_loadable_modules:
-                self.assertTrue(check_if_lora_correctly_set(pipe.text_encoder), "Unfuse should still keep LoRA layers")
-
-            self.assertTrue(check_if_lora_correctly_set(denoiser), "Unfuse should still keep LoRA layers")
-
-            if self.has_two_text_encoders or self.has_three_text_encoders:
-                if "text_encoder_2" in self.pipeline_class._lora_loadable_modules:
-                    self.assertTrue(
-                        check_if_lora_correctly_set(pipe.text_encoder_2), "Unfuse should still keep LoRA layers"
-                    )
-
             pipe.fuse_lora(
                 components=self.pipeline_class._lora_loadable_modules, adapter_names=["adapter-2", "adapter-1"]
             )
-            self.assertTrue(pipe.num_fused_loras == 2, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
 
             # Fusing should still keep the LoRA layers
             output_all_lora_fused = pipe(**inputs, generator=torch.manual_seed(0))[0]
@@ -1669,63 +1657,6 @@ class PeftLoraLoaderMixinTests:
                 np.allclose(output_all_lora_fused, outputs_all_lora, atol=expected_atol, rtol=expected_rtol),
                 "Fused lora should not change the output",
             )
-            pipe.unfuse_lora(components=self.pipeline_class._lora_loadable_modules)
-            self.assertTrue(pipe.num_fused_loras == 0, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
-
-    def test_lora_scale_kwargs_match_fusion(self, expected_atol: float = 1e-3, expected_rtol: float = 1e-3):
-        attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class)
-
-        for lora_scale in [1.0, 0.8]:
-            for scheduler_cls in self.scheduler_classes:
-                components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls)
-                pipe = self.pipeline_class(**components)
-                pipe = pipe.to(torch_device)
-                pipe.set_progress_bar_config(disable=None)
-                _, _, inputs = self.get_dummy_inputs(with_generator=False)
-
-                output_no_lora = pipe(**inputs, generator=torch.manual_seed(0))[0]
-                self.assertTrue(output_no_lora.shape == self.output_shape)
-
-                if "text_encoder" in self.pipeline_class._lora_loadable_modules:
-                    pipe.text_encoder.add_adapter(text_lora_config, "adapter-1")
-                    self.assertTrue(
-                        check_if_lora_correctly_set(pipe.text_encoder), "Lora not correctly set in text encoder"
-                    )
-
-                denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet
-                denoiser.add_adapter(denoiser_lora_config, "adapter-1")
-                self.assertTrue(check_if_lora_correctly_set(denoiser), "Lora not correctly set in denoiser.")
-
-                if self.has_two_text_encoders or self.has_three_text_encoders:
-                    lora_loadable_components = self.pipeline_class._lora_loadable_modules
-                    if "text_encoder_2" in lora_loadable_components:
-                        pipe.text_encoder_2.add_adapter(text_lora_config, "adapter-1")
-                        self.assertTrue(
-                            check_if_lora_correctly_set(pipe.text_encoder_2),
-                            "Lora not correctly set in text encoder 2",
-                        )
-
-                pipe.set_adapters(["adapter-1"])
-                attention_kwargs = {attention_kwargs_name: {"scale": lora_scale}}
-                outputs_lora_1 = pipe(**inputs, generator=torch.manual_seed(0), **attention_kwargs)[0]
-
-                pipe.fuse_lora(
-                    components=self.pipeline_class._lora_loadable_modules,
-                    adapter_names=["adapter-1"],
-                    lora_scale=lora_scale,
-                )
-                self.assertTrue(pipe.num_fused_loras == 1, f"{pipe.num_fused_loras=}, {pipe.fused_loras=}")
-
-                outputs_lora_1_fused = pipe(**inputs, generator=torch.manual_seed(0))[0]
-
-                self.assertTrue(
-                    np.allclose(outputs_lora_1, outputs_lora_1_fused, atol=expected_atol, rtol=expected_rtol),
-                    "Fused lora should not change the output",
-                )
-                self.assertFalse(
-                    np.allclose(output_no_lora, outputs_lora_1, atol=expected_atol, rtol=expected_rtol),
-                    "LoRA should change the output",
-                )
 
     @require_peft_version_greater(peft_version="0.9.0")
     def test_simple_inference_with_dora(self):
@@ -1907,7 +1838,12 @@ class PeftLoraLoaderMixinTests:
 
     def test_set_adapters_match_attention_kwargs(self):
         """Test to check if outputs after `set_adapters()` and attention kwargs match."""
-        attention_kwargs_name = determine_attention_kwargs_name(self.pipeline_class)
+        call_signature_keys = inspect.signature(self.pipeline_class.__call__).parameters.keys()
+        for possible_attention_kwargs in POSSIBLE_ATTENTION_KWARGS_NAMES:
+            if possible_attention_kwargs in call_signature_keys:
+                attention_kwargs_name = possible_attention_kwargs
+                break
+        assert attention_kwargs_name is not None
 
         for scheduler_cls in self.scheduler_classes:
             components, text_lora_config, denoiser_lora_config = self.get_dummy_components(scheduler_cls)

tests/models/test_modeling_common.py

@@ -1060,7 +1060,7 @@ class ModelTesterMixin:
     @parameterized.expand([True, False])
     @torch.no_grad()
     @unittest.skipIf(not is_peft_available(), "Only with PEFT")
-    def test_lora_save_load_adapter(self, use_dora=False):
+    def test_save_load_lora_adapter(self, use_dora=False):
         import safetensors
         from peft import LoraConfig
         from peft.utils import get_peft_model_state_dict
@@ -1117,7 +1117,7 @@ class ModelTesterMixin:
         self.assertTrue(torch.allclose(outputs_with_lora, outputs_with_lora_2, atol=1e-4, rtol=1e-4))
 
     @unittest.skipIf(not is_peft_available(), "Only with PEFT")
-    def test_lora_wrong_adapter_name_raises_error(self):
+    def test_wrong_adapter_name_raises_error(self):
         from peft import LoraConfig
 
         from diffusers.loaders.peft import PeftAdapterMixin
@@ -1744,10 +1744,6 @@ class ModelPushToHubTester(unittest.TestCase):
         delete_repo(self.repo_id, token=TOKEN)
 
 
-@require_torch_gpu
-@require_torch_2
-@is_torch_compile
-@slow
 class TorchCompileTesterMixin:
     def setUp(self):
         # clean up the VRAM before each test
@@ -1763,7 +1759,12 @@ class TorchCompileTesterMixin:
         gc.collect()
         backend_empty_cache(torch_device)
 
+    @require_torch_gpu
+    @require_torch_2
+    @is_torch_compile
+    @slow
     def test_torch_compile_recompilation_and_graph_break(self):
+        torch.compiler.reset()
         init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
 
         model = self.model_class(**init_dict).to(torch_device)
@@ -1777,31 +1778,6 @@ class TorchCompileTesterMixin:
             _ = model(**inputs_dict)
             _ = model(**inputs_dict)
 
-    def test_compile_with_group_offloading(self):
-        torch._dynamo.config.cache_size_limit = 10000
-
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
-
-        if not getattr(model, "_supports_group_offloading", True):
-            return
-
-        model.eval()
-        # TODO: Can test for other group offloading kwargs later if needed.
-        group_offload_kwargs = {
-            "onload_device": "cuda",
-            "offload_device": "cpu",
-            "offload_type": "block_level",
-            "num_blocks_per_group": 1,
-            "use_stream": True,
-            "non_blocking": True,
-        }
-        model.enable_group_offload(**group_offload_kwargs)
-        model.compile()
-        with torch.no_grad():
-            _ = model(**inputs_dict)
-            _ = model(**inputs_dict)
-
 
 @slow
 @require_torch_2

tests/pipelines/cogvideo/test_cogvideox_image2video.py

@@ -270,7 +270,7 @@ class CogVideoXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestC
         generator_device = "cpu"
         components = self.get_dummy_components()
 
-        # The reason to modify it this way is because I2V Transformer limits the generation to resolutions used during initialization.
+        # The reason to modify it this way is because I2V Transformer limits the generation to resolutions used during initalization.
         # This limitation comes from using learned positional embeddings which cannot be generated on-the-fly like sincos or RoPE embeddings.
         # See the if-statement on "self.use_learned_positional_embeddings" in diffusers/models/embeddings.py
         components["transformer"] = CogVideoXTransformer3DModel.from_config(

tests/pipelines/consisid/test_consisid.py

@@ -280,7 +280,7 @@ class ConsisIDPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
         generator_device = "cpu"
         components = self.get_dummy_components()
 
-        # The reason to modify it this way is because ConsisID Transformer limits the generation to resolutions used during initialization.
+        # The reason to modify it this way is because ConsisID Transformer limits the generation to resolutions used during initalization.
         # This limitation comes from using learned positional embeddings which cannot be generated on-the-fly like sincos or RoPE embeddings.
         # See the if-statement on "self.use_learned_positional_embeddings" in diffusers/models/embeddings.py
         components["transformer"] = ConsisIDTransformer3DModel.from_config(

tests/pipelines/cosmos/test_cosmos2_text2_image.py

@@ -1,342 +0,0 @@
-# Copyright 2024 The HuggingFace Team.
-#
-# 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 inspect
-import json
-import os
-import tempfile
-import unittest
-
-import numpy as np
-import torch
-from transformers import AutoTokenizer, T5EncoderModel
-
-from diffusers import AutoencoderKLWan, Cosmos2TextToImagePipeline, CosmosTransformer3DModel, EDMEulerScheduler
-from diffusers.utils.testing_utils import enable_full_determinism, torch_device
-
-from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import PipelineTesterMixin, to_np
-from .cosmos_guardrail import DummyCosmosSafetyChecker
-
-
-enable_full_determinism()
-
-
-class Cosmos2TextToImagePipelineWrapper(Cosmos2TextToImagePipeline):
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        kwargs["safety_checker"] = DummyCosmosSafetyChecker()
-        return Cosmos2TextToImagePipeline.from_pretrained(*args, **kwargs)
-
-
-class Cosmos2TextToImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = Cosmos2TextToImagePipelineWrapper
-    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
-    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
-    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    required_optional_params = frozenset(
-        [
-            "num_inference_steps",
-            "generator",
-            "latents",
-            "return_dict",
-            "callback_on_step_end",
-            "callback_on_step_end_tensor_inputs",
-        ]
-    )
-    supports_dduf = False
-    test_xformers_attention = False
-    test_layerwise_casting = True
-    test_group_offloading = True
-
-    def get_dummy_components(self):
-        torch.manual_seed(0)
-        transformer = CosmosTransformer3DModel(
-            in_channels=16,
-            out_channels=16,
-            num_attention_heads=2,
-            attention_head_dim=16,
-            num_layers=2,
-            mlp_ratio=2,
-            text_embed_dim=32,
-            adaln_lora_dim=4,
-            max_size=(4, 32, 32),
-            patch_size=(1, 2, 2),
-            rope_scale=(2.0, 1.0, 1.0),
-            concat_padding_mask=True,
-            extra_pos_embed_type="learnable",
-        )
-
-        torch.manual_seed(0)
-        vae = AutoencoderKLWan(
-            base_dim=3,
-            z_dim=16,
-            dim_mult=[1, 1, 1, 1],
-            num_res_blocks=1,
-            temperal_downsample=[False, True, True],
-        )
-
-        torch.manual_seed(0)
-        scheduler = EDMEulerScheduler(
-            sigma_min=0.002,
-            sigma_max=80,
-            sigma_data=0.5,
-            sigma_schedule="karras",
-            num_train_timesteps=1000,
-            prediction_type="epsilon",
-            rho=7.0,
-            final_sigmas_type="sigma_min",
-            use_flow_sigmas=True,
-        )
-        text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
-        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
-
-        components = {
-            "transformer": transformer,
-            "vae": vae,
-            "scheduler": scheduler,
-            "text_encoder": text_encoder,
-            "tokenizer": tokenizer,
-            # We cannot run the Cosmos Guardrail for fast tests due to the large model size
-            "safety_checker": DummyCosmosSafetyChecker(),
-        }
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-
-        inputs = {
-            "prompt": "dance monkey",
-            "negative_prompt": "bad quality",
-            "generator": generator,
-            "num_inference_steps": 2,
-            "guidance_scale": 3.0,
-            "height": 32,
-            "width": 32,
-            "max_sequence_length": 16,
-            "output_type": "pt",
-        }
-
-        return inputs
-
-    def test_inference(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        image = pipe(**inputs).images
-        generated_image = image[0]
-
-        self.assertEqual(generated_image.shape, (3, 32, 32))
-        expected_video = torch.randn(3, 32, 32)
-        max_diff = np.abs(generated_image - expected_video).max()
-        self.assertLessEqual(max_diff, 1e10)
-
-    def test_callback_inputs(self):
-        sig = inspect.signature(self.pipeline_class.__call__)
-        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
-        has_callback_step_end = "callback_on_step_end" in sig.parameters
-
-        if not (has_callback_tensor_inputs and has_callback_step_end):
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe = pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-        self.assertTrue(
-            hasattr(pipe, "_callback_tensor_inputs"),
-            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
-        )
-
-        def callback_inputs_subset(pipe, i, t, callback_kwargs):
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        def callback_inputs_all(pipe, i, t, callback_kwargs):
-            for tensor_name in pipe._callback_tensor_inputs:
-                assert tensor_name in callback_kwargs
-
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        inputs = self.get_dummy_inputs(torch_device)
-
-        # Test passing in a subset
-        inputs["callback_on_step_end"] = callback_inputs_subset
-        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
-        output = pipe(**inputs)[0]
-
-        # Test passing in a everything
-        inputs["callback_on_step_end"] = callback_inputs_all
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-
-        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
-            is_last = i == (pipe.num_timesteps - 1)
-            if is_last:
-                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
-            return callback_kwargs
-
-        inputs["callback_on_step_end"] = callback_inputs_change_tensor
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-        assert output.abs().sum() < 1e10
-
-    def test_inference_batch_single_identical(self):
-        self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-2)
-
-    def test_attention_slicing_forward_pass(
-        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
-    ):
-        if not self.test_attention_slicing:
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        for component in pipe.components.values():
-            if hasattr(component, "set_default_attn_processor"):
-                component.set_default_attn_processor()
-        pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-
-        generator_device = "cpu"
-        inputs = self.get_dummy_inputs(generator_device)
-        output_without_slicing = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=1)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing1 = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=2)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing2 = pipe(**inputs)[0]
-
-        if test_max_difference:
-            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
-            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
-            self.assertLess(
-                max(max_diff1, max_diff2),
-                expected_max_diff,
-                "Attention slicing should not affect the inference results",
-            )
-
-    def test_vae_tiling(self, expected_diff_max: float = 0.2):
-        generator_device = "cpu"
-        components = self.get_dummy_components()
-
-        pipe = self.pipeline_class(**components)
-        pipe.to("cpu")
-        pipe.set_progress_bar_config(disable=None)
-
-        # Without tiling
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_without_tiling = pipe(**inputs)[0]
-
-        # With tiling
-        pipe.vae.enable_tiling(
-            tile_sample_min_height=96,
-            tile_sample_min_width=96,
-            tile_sample_stride_height=64,
-            tile_sample_stride_width=64,
-        )
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_with_tiling = pipe(**inputs)[0]
-
-        self.assertLess(
-            (to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
-            expected_diff_max,
-            "VAE tiling should not affect the inference results",
-        )
-
-    def test_save_load_optional_components(self, expected_max_difference=1e-4):
-        self.pipeline_class._optional_components.remove("safety_checker")
-        super().test_save_load_optional_components(expected_max_difference=expected_max_difference)
-        self.pipeline_class._optional_components.append("safety_checker")
-
-    def test_serialization_with_variants(self):
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        model_components = [
-            component_name
-            for component_name, component in pipe.components.items()
-            if isinstance(component, torch.nn.Module)
-        ]
-        model_components.remove("safety_checker")
-        variant = "fp16"
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False)
-
-            with open(f"{tmpdir}/model_index.json", "r") as f:
-                config = json.load(f)
-
-            for subfolder in os.listdir(tmpdir):
-                if not os.path.isfile(subfolder) and subfolder in model_components:
-                    folder_path = os.path.join(tmpdir, subfolder)
-                    is_folder = os.path.isdir(folder_path) and subfolder in config
-                    assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path))
-
-    def test_torch_dtype_dict(self):
-        components = self.get_dummy_components()
-        if not components:
-            self.skipTest("No dummy components defined.")
-
-        pipe = self.pipeline_class(**components)
-
-        specified_key = next(iter(components.keys()))
-
-        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
-            pipe.save_pretrained(tmpdirname, safe_serialization=False)
-            torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16}
-            loaded_pipe = self.pipeline_class.from_pretrained(
-                tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict
-            )
-
-        for name, component in loaded_pipe.components.items():
-            if name == "safety_checker":
-                continue
-            if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"):
-                expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32))
-                self.assertEqual(
-                    component.dtype,
-                    expected_dtype,
-                    f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}",
-                )
-
-    @unittest.skip(
-        "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in "
-        "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is "
-        "too large and slow to run on CI."
-    )
-    def test_encode_prompt_works_in_isolation(self):
-        pass

tests/pipelines/cosmos/test_cosmos2_video2world.py

@@ -1,356 +0,0 @@
-# Copyright 2024 The HuggingFace Team.
-#
-# 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 inspect
-import json
-import os
-import tempfile
-import unittest
-
-import numpy as np
-import PIL.Image
-import torch
-from transformers import AutoTokenizer, T5EncoderModel
-
-from diffusers import AutoencoderKLWan, Cosmos2VideoToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler
-from diffusers.utils.testing_utils import enable_full_determinism, torch_device
-
-from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import PipelineTesterMixin, to_np
-from .cosmos_guardrail import DummyCosmosSafetyChecker
-
-
-enable_full_determinism()
-
-
-class Cosmos2VideoToWorldPipelineWrapper(Cosmos2VideoToWorldPipeline):
-    @staticmethod
-    def from_pretrained(*args, **kwargs):
-        kwargs["safety_checker"] = DummyCosmosSafetyChecker()
-        return Cosmos2VideoToWorldPipeline.from_pretrained(*args, **kwargs)
-
-
-class Cosmos2VideoToWorldPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = Cosmos2VideoToWorldPipelineWrapper
-    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
-    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image", "video"})
-    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    required_optional_params = frozenset(
-        [
-            "num_inference_steps",
-            "generator",
-            "latents",
-            "return_dict",
-            "callback_on_step_end",
-            "callback_on_step_end_tensor_inputs",
-        ]
-    )
-    supports_dduf = False
-    test_xformers_attention = False
-    test_layerwise_casting = True
-    test_group_offloading = True
-
-    def get_dummy_components(self):
-        torch.manual_seed(0)
-        transformer = CosmosTransformer3DModel(
-            in_channels=16 + 1,
-            out_channels=16,
-            num_attention_heads=2,
-            attention_head_dim=16,
-            num_layers=2,
-            mlp_ratio=2,
-            text_embed_dim=32,
-            adaln_lora_dim=4,
-            max_size=(4, 32, 32),
-            patch_size=(1, 2, 2),
-            rope_scale=(2.0, 1.0, 1.0),
-            concat_padding_mask=True,
-            extra_pos_embed_type="learnable",
-        )
-
-        torch.manual_seed(0)
-        vae = AutoencoderKLWan(
-            base_dim=3,
-            z_dim=16,
-            dim_mult=[1, 1, 1, 1],
-            num_res_blocks=1,
-            temperal_downsample=[False, True, True],
-        )
-
-        torch.manual_seed(0)
-        scheduler = EDMEulerScheduler(
-            sigma_min=0.002,
-            sigma_max=80,
-            sigma_data=0.5,
-            sigma_schedule="karras",
-            num_train_timesteps=1000,
-            prediction_type="epsilon",
-            rho=7.0,
-            final_sigmas_type="sigma_min",
-            use_flow_sigmas=True,
-        )
-        text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
-        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
-
-        components = {
-            "transformer": transformer,
-            "vae": vae,
-            "scheduler": scheduler,
-            "text_encoder": text_encoder,
-            "tokenizer": tokenizer,
-            # We cannot run the Cosmos Guardrail for fast tests due to the large model size
-            "safety_checker": DummyCosmosSafetyChecker(),
-        }
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-
-        image_height = 32
-        image_width = 32
-        image = PIL.Image.new("RGB", (image_width, image_height))
-
-        inputs = {
-            "image": image,
-            "prompt": "dance monkey",
-            "negative_prompt": "bad quality",
-            "generator": generator,
-            "num_inference_steps": 2,
-            "guidance_scale": 3.0,
-            "height": image_height,
-            "width": image_width,
-            "num_frames": 9,
-            "max_sequence_length": 16,
-            "output_type": "pt",
-        }
-
-        return inputs
-
-    def test_inference(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        video = pipe(**inputs).frames
-        generated_video = video[0]
-
-        self.assertEqual(generated_video.shape, (9, 3, 32, 32))
-        expected_video = torch.randn(9, 3, 32, 32)
-        max_diff = np.abs(generated_video - expected_video).max()
-        self.assertLessEqual(max_diff, 1e10)
-
-    def test_components_function(self):
-        init_components = self.get_dummy_components()
-        init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))}
-        pipe = self.pipeline_class(**init_components)
-        self.assertTrue(hasattr(pipe, "components"))
-        self.assertTrue(set(pipe.components.keys()) == set(init_components.keys()))
-
-    def test_callback_inputs(self):
-        sig = inspect.signature(self.pipeline_class.__call__)
-        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
-        has_callback_step_end = "callback_on_step_end" in sig.parameters
-
-        if not (has_callback_tensor_inputs and has_callback_step_end):
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe = pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-        self.assertTrue(
-            hasattr(pipe, "_callback_tensor_inputs"),
-            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
-        )
-
-        def callback_inputs_subset(pipe, i, t, callback_kwargs):
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        def callback_inputs_all(pipe, i, t, callback_kwargs):
-            for tensor_name in pipe._callback_tensor_inputs:
-                assert tensor_name in callback_kwargs
-
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        inputs = self.get_dummy_inputs(torch_device)
-
-        # Test passing in a subset
-        inputs["callback_on_step_end"] = callback_inputs_subset
-        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
-        output = pipe(**inputs)[0]
-
-        # Test passing in a everything
-        inputs["callback_on_step_end"] = callback_inputs_all
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-
-        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
-            is_last = i == (pipe.num_timesteps - 1)
-            if is_last:
-                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
-            return callback_kwargs
-
-        inputs["callback_on_step_end"] = callback_inputs_change_tensor
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-        assert output.abs().sum() < 1e10
-
-    def test_inference_batch_single_identical(self):
-        self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-2)
-
-    def test_attention_slicing_forward_pass(
-        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
-    ):
-        if not self.test_attention_slicing:
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        for component in pipe.components.values():
-            if hasattr(component, "set_default_attn_processor"):
-                component.set_default_attn_processor()
-        pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-
-        generator_device = "cpu"
-        inputs = self.get_dummy_inputs(generator_device)
-        output_without_slicing = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=1)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing1 = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=2)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing2 = pipe(**inputs)[0]
-
-        if test_max_difference:
-            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
-            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
-            self.assertLess(
-                max(max_diff1, max_diff2),
-                expected_max_diff,
-                "Attention slicing should not affect the inference results",
-            )
-
-    def test_vae_tiling(self, expected_diff_max: float = 0.2):
-        generator_device = "cpu"
-        components = self.get_dummy_components()
-
-        pipe = self.pipeline_class(**components)
-        pipe.to("cpu")
-        pipe.set_progress_bar_config(disable=None)
-
-        # Without tiling
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_without_tiling = pipe(**inputs)[0]
-
-        # With tiling
-        pipe.vae.enable_tiling(
-            tile_sample_min_height=96,
-            tile_sample_min_width=96,
-            tile_sample_stride_height=64,
-            tile_sample_stride_width=64,
-        )
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_with_tiling = pipe(**inputs)[0]
-
-        self.assertLess(
-            (to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
-            expected_diff_max,
-            "VAE tiling should not affect the inference results",
-        )
-
-    def test_save_load_optional_components(self, expected_max_difference=1e-4):
-        self.pipeline_class._optional_components.remove("safety_checker")
-        super().test_save_load_optional_components(expected_max_difference=expected_max_difference)
-        self.pipeline_class._optional_components.append("safety_checker")
-
-    def test_serialization_with_variants(self):
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        model_components = [
-            component_name
-            for component_name, component in pipe.components.items()
-            if isinstance(component, torch.nn.Module)
-        ]
-        model_components.remove("safety_checker")
-        variant = "fp16"
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False)
-
-            with open(f"{tmpdir}/model_index.json", "r") as f:
-                config = json.load(f)
-
-            for subfolder in os.listdir(tmpdir):
-                if not os.path.isfile(subfolder) and subfolder in model_components:
-                    folder_path = os.path.join(tmpdir, subfolder)
-                    is_folder = os.path.isdir(folder_path) and subfolder in config
-                    assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path))
-
-    def test_torch_dtype_dict(self):
-        components = self.get_dummy_components()
-        if not components:
-            self.skipTest("No dummy components defined.")
-
-        pipe = self.pipeline_class(**components)
-
-        specified_key = next(iter(components.keys()))
-
-        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
-            pipe.save_pretrained(tmpdirname, safe_serialization=False)
-            torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16}
-            loaded_pipe = self.pipeline_class.from_pretrained(
-                tmpdirname, safety_checker=DummyCosmosSafetyChecker(), torch_dtype=torch_dtype_dict
-            )
-
-        for name, component in loaded_pipe.components.items():
-            if name == "safety_checker":
-                continue
-            if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"):
-                expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32))
-                self.assertEqual(
-                    component.dtype,
-                    expected_dtype,
-                    f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}",
-                )
-
-    @unittest.skip(
-        "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in "
-        "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is "
-        "too large and slow to run on CI."
-    )
-    def test_encode_prompt_works_in_isolation(self):
-        pass

tests/pipelines/kolors/test_kolors_img2img.py

@@ -155,6 +155,6 @@ class KolorsPipelineImg2ImgFastTests(PipelineTesterMixin, unittest.TestCase):
     def test_float16_inference(self):
         super().test_float16_inference(expected_max_diff=7e-2)
 
-    @unittest.skip("Test not supported because kolors img2img doesn't take pooled embeds as inputs unlike kolors t2i.")
+    @unittest.skip("Test not supported because kolors img2img doesn't take pooled embeds as inputs unline kolors t2i.")
     def test_encode_prompt_works_in_isolation(self):
         pass

tests/pipelines/pag/test_pag_pixart_sigma.py

@@ -254,7 +254,7 @@ class PixArtSigmaPAGPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
             assert_mean_pixel_difference(to_np(output_with_slicing1[0]), to_np(output_without_slicing[0]))
             assert_mean_pixel_difference(to_np(output_with_slicing2[0]), to_np(output_without_slicing[0]))
 
-    # Because we have `pag_applied_layers` we cannot directly apply
+    # Because we have `pag_applied_layers` we cannot direcly apply
     # `set_default_attn_processor`
     def test_dict_tuple_outputs_equivalent(self, expected_slice=None, expected_max_difference=1e-4):
         components = self.get_dummy_components()

tests/pipelines/sana/test_sana_sprint_img2img.py

@@ -1,313 +0,0 @@
-# Copyright 2024 The HuggingFace Team.
-#
-# 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 inspect
-import unittest
-
-import numpy as np
-import torch
-from transformers import Gemma2Config, Gemma2Model, GemmaTokenizer
-
-from diffusers import AutoencoderDC, SanaSprintImg2ImgPipeline, SanaTransformer2DModel, SCMScheduler
-from diffusers.utils.testing_utils import (
-    enable_full_determinism,
-    torch_device,
-)
-
-from ..pipeline_params import (
-    IMAGE_TO_IMAGE_IMAGE_PARAMS,
-    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
-    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
-)
-from ..test_pipelines_common import PipelineTesterMixin, to_np
-
-
-enable_full_determinism()
-
-
-class SanaSprintImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = SanaSprintImg2ImgPipeline
-    params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
-        "negative_prompt",
-        "negative_prompt_embeds",
-    }
-    batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - {"negative_prompt"}
-    image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
-    image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
-    required_optional_params = frozenset(
-        [
-            "num_inference_steps",
-            "generator",
-            "latents",
-            "return_dict",
-            "callback_on_step_end",
-            "callback_on_step_end_tensor_inputs",
-        ]
-    )
-    test_xformers_attention = False
-    test_layerwise_casting = True
-    test_group_offloading = True
-
-    def get_dummy_components(self):
-        torch.manual_seed(0)
-        transformer = SanaTransformer2DModel(
-            patch_size=1,
-            in_channels=4,
-            out_channels=4,
-            num_layers=1,
-            num_attention_heads=2,
-            attention_head_dim=4,
-            num_cross_attention_heads=2,
-            cross_attention_head_dim=4,
-            cross_attention_dim=8,
-            caption_channels=8,
-            sample_size=32,
-            qk_norm="rms_norm_across_heads",
-            guidance_embeds=True,
-        )
-
-        torch.manual_seed(0)
-        vae = AutoencoderDC(
-            in_channels=3,
-            latent_channels=4,
-            attention_head_dim=2,
-            encoder_block_types=(
-                "ResBlock",
-                "EfficientViTBlock",
-            ),
-            decoder_block_types=(
-                "ResBlock",
-                "EfficientViTBlock",
-            ),
-            encoder_block_out_channels=(8, 8),
-            decoder_block_out_channels=(8, 8),
-            encoder_qkv_multiscales=((), (5,)),
-            decoder_qkv_multiscales=((), (5,)),
-            encoder_layers_per_block=(1, 1),
-            decoder_layers_per_block=[1, 1],
-            downsample_block_type="conv",
-            upsample_block_type="interpolate",
-            decoder_norm_types="rms_norm",
-            decoder_act_fns="silu",
-            scaling_factor=0.41407,
-        )
-
-        torch.manual_seed(0)
-        scheduler = SCMScheduler()
-
-        torch.manual_seed(0)
-        text_encoder_config = Gemma2Config(
-            head_dim=16,
-            hidden_size=8,
-            initializer_range=0.02,
-            intermediate_size=64,
-            max_position_embeddings=8192,
-            model_type="gemma2",
-            num_attention_heads=2,
-            num_hidden_layers=1,
-            num_key_value_heads=2,
-            vocab_size=8,
-            attn_implementation="eager",
-        )
-        text_encoder = Gemma2Model(text_encoder_config)
-        tokenizer = GemmaTokenizer.from_pretrained("hf-internal-testing/dummy-gemma")
-
-        components = {
-            "transformer": transformer,
-            "vae": vae,
-            "scheduler": scheduler,
-            "text_encoder": text_encoder,
-            "tokenizer": tokenizer,
-        }
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-        image = torch.randn(1, 3, 32, 32, generator=generator)
-        inputs = {
-            "prompt": "",
-            "image": image,
-            "strength": 0.5,
-            "generator": generator,
-            "num_inference_steps": 2,
-            "guidance_scale": 6.0,
-            "height": 32,
-            "width": 32,
-            "max_sequence_length": 16,
-            "output_type": "pt",
-            "complex_human_instruction": None,
-        }
-        return inputs
-
-    def test_inference(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        image = pipe(**inputs)[0]
-        generated_image = image[0]
-
-        self.assertEqual(generated_image.shape, (3, 32, 32))
-        expected_image = torch.randn(3, 32, 32)
-        max_diff = np.abs(generated_image - expected_image).max()
-        self.assertLessEqual(max_diff, 1e10)
-
-    def test_callback_inputs(self):
-        sig = inspect.signature(self.pipeline_class.__call__)
-        has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
-        has_callback_step_end = "callback_on_step_end" in sig.parameters
-
-        if not (has_callback_tensor_inputs and has_callback_step_end):
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe = pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-        self.assertTrue(
-            hasattr(pipe, "_callback_tensor_inputs"),
-            f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
-        )
-
-        def callback_inputs_subset(pipe, i, t, callback_kwargs):
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        def callback_inputs_all(pipe, i, t, callback_kwargs):
-            for tensor_name in pipe._callback_tensor_inputs:
-                assert tensor_name in callback_kwargs
-
-            # iterate over callback args
-            for tensor_name, tensor_value in callback_kwargs.items():
-                # check that we're only passing in allowed tensor inputs
-                assert tensor_name in pipe._callback_tensor_inputs
-
-            return callback_kwargs
-
-        inputs = self.get_dummy_inputs(torch_device)
-
-        # Test passing in a subset
-        inputs["callback_on_step_end"] = callback_inputs_subset
-        inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
-        output = pipe(**inputs)[0]
-
-        # Test passing in a everything
-        inputs["callback_on_step_end"] = callback_inputs_all
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-
-        def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
-            is_last = i == (pipe.num_timesteps - 1)
-            if is_last:
-                callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
-            return callback_kwargs
-
-        inputs["callback_on_step_end"] = callback_inputs_change_tensor
-        inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
-        output = pipe(**inputs)[0]
-        assert output.abs().sum() < 1e10
-
-    def test_attention_slicing_forward_pass(
-        self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
-    ):
-        if not self.test_attention_slicing:
-            return
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        for component in pipe.components.values():
-            if hasattr(component, "set_default_attn_processor"):
-                component.set_default_attn_processor()
-        pipe.to(torch_device)
-        pipe.set_progress_bar_config(disable=None)
-
-        generator_device = "cpu"
-        inputs = self.get_dummy_inputs(generator_device)
-        output_without_slicing = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=1)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing1 = pipe(**inputs)[0]
-
-        pipe.enable_attention_slicing(slice_size=2)
-        inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing2 = pipe(**inputs)[0]
-
-        if test_max_difference:
-            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
-            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
-            self.assertLess(
-                max(max_diff1, max_diff2),
-                expected_max_diff,
-                "Attention slicing should not affect the inference results",
-            )
-
-    @unittest.skip("vae tiling resulted in a small margin over the expected max diff, so skipping this test for now")
-    def test_vae_tiling(self, expected_diff_max: float = 0.2):
-        generator_device = "cpu"
-        components = self.get_dummy_components()
-
-        pipe = self.pipeline_class(**components)
-        pipe.to("cpu")
-        pipe.set_progress_bar_config(disable=None)
-
-        # Without tiling
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_without_tiling = pipe(**inputs)[0]
-
-        # With tiling
-        pipe.vae.enable_tiling(
-            tile_sample_min_height=96,
-            tile_sample_min_width=96,
-            tile_sample_stride_height=64,
-            tile_sample_stride_width=64,
-        )
-        inputs = self.get_dummy_inputs(generator_device)
-        inputs["height"] = inputs["width"] = 128
-        output_with_tiling = pipe(**inputs)[0]
-
-        self.assertLess(
-            (to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
-            expected_diff_max,
-            "VAE tiling should not affect the inference results",
-        )
-
-    # TODO(aryan): Create a dummy gemma model with smol vocab size
-    @unittest.skip(
-        "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
-    )
-    def test_inference_batch_consistent(self):
-        pass
-
-    @unittest.skip(
-        "A very small vocab size is used for fast tests. So, any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error."
-    )
-    def test_inference_batch_single_identical(self):
-        pass
-
-    def test_float16_inference(self):
-        # Requires higher tolerance as model seems very sensitive to dtype
-        super().test_float16_inference(expected_max_diff=0.08)

tests/pipelines/stable_unclip/test_stable_unclip.py

@@ -227,7 +227,7 @@ class StableUnCLIPPipelineIntegrationTests(unittest.TestCase):
         pipe.enable_sequential_cpu_offload()
 
         generator = torch.Generator(device="cpu").manual_seed(0)
-        output = pipe("anime turtle", generator=generator, output_type="np")
+        output = pipe("anime turle", generator=generator, output_type="np")
 
         image = output.images[0]
 

tests/pipelines/stable_unclip/test_stable_unclip_img2img.py

@@ -250,7 +250,7 @@ class StableUnCLIPImg2ImgPipelineIntegrationTests(unittest.TestCase):
         pipe.enable_sequential_cpu_offload()
 
         generator = torch.Generator(device="cpu").manual_seed(0)
-        output = pipe(input_image, "anime turtle", generator=generator, output_type="np")
+        output = pipe(input_image, "anime turle", generator=generator, output_type="np")
 
         image = output.images[0]
 
@@ -277,7 +277,7 @@ class StableUnCLIPImg2ImgPipelineIntegrationTests(unittest.TestCase):
         pipe.enable_sequential_cpu_offload()
 
         generator = torch.Generator(device="cpu").manual_seed(0)
-        output = pipe(input_image, "anime turtle", generator=generator, output_type="np")
+        output = pipe(input_image, "anime turle", generator=generator, output_type="np")
 
         image = output.images[0]
 

tests/pipelines/test_pipeline_utils.py

@@ -19,7 +19,7 @@ from diffusers import (
     UNet2DConditionModel,
 )
 from diffusers.pipelines.pipeline_loading_utils import is_safetensors_compatible, variant_compatible_siblings
-from diffusers.utils.testing_utils import require_torch_accelerator, torch_device
+from diffusers.utils.testing_utils import require_torch_gpu, torch_device
 
 
 class IsSafetensorsCompatibleTests(unittest.TestCase):
@@ -850,9 +850,9 @@ class ProgressBarTests(unittest.TestCase):
             self.assertTrue(stderr.getvalue() == "", "Progress bar should be disabled")
 
 
-@require_torch_accelerator
+@require_torch_gpu
 class PipelineDeviceAndDtypeStabilityTests(unittest.TestCase):
-    expected_pipe_device = torch.device(f"{torch_device}:0")
+    expected_pipe_device = torch.device("cuda:0")
     expected_pipe_dtype = torch.float64
 
     def get_dummy_components_image_generation(self):
@@ -921,8 +921,8 @@ class PipelineDeviceAndDtypeStabilityTests(unittest.TestCase):
         pipe.to(device=torch_device, dtype=torch.float32)
 
         pipe.unet.to(device="cpu")
-        pipe.vae.to(device=torch_device)
-        pipe.text_encoder.to(device=f"{torch_device}:0")
+        pipe.vae.to(device="cuda")
+        pipe.text_encoder.to(device="cuda:0")
 
         pipe_device = pipe.device
 

tests/pipelines/test_pipelines_common.py

@@ -2096,11 +2096,11 @@ class PipelineTesterMixin:
         with torch.no_grad():
             encoded_prompt_outputs = pipe_with_just_text_encoder.encode_prompt(**encode_prompt_inputs)
 
-        # Programmatically determine the return names of `encode_prompt.`
-        ast_visitor = ReturnNameVisitor()
-        encode_prompt_tree = ast_visitor.get_ast_tree(cls=self.pipeline_class)
-        ast_visitor.visit(encode_prompt_tree)
-        prompt_embed_kwargs = ast_visitor.return_names
+        # Programatically determine the reutrn names of `encode_prompt.`
+        ast_vistor = ReturnNameVisitor()
+        encode_prompt_tree = ast_vistor.get_ast_tree(cls=self.pipeline_class)
+        ast_vistor.visit(encode_prompt_tree)
+        prompt_embed_kwargs = ast_vistor.return_names
         prompt_embeds_kwargs = dict(zip(prompt_embed_kwargs, encoded_prompt_outputs))
 
         # Pack the outputs of `encode_prompt`.

tests/pipelines/wan/test_wan_vace.py

@@ -1,201 +0,0 @@
-# Copyright 2025 The HuggingFace Team.
-#
-# 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 unittest
-
-import numpy as np
-import torch
-from PIL import Image
-from transformers import AutoTokenizer, T5EncoderModel
-
-from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler, WanVACEPipeline, WanVACETransformer3DModel
-from diffusers.utils.testing_utils import enable_full_determinism
-
-from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
-from ..test_pipelines_common import PipelineTesterMixin
-
-
-enable_full_determinism()
-
-
-class WanVACEPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
-    pipeline_class = WanVACEPipeline
-    params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
-    batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
-    image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
-    required_optional_params = frozenset(
-        [
-            "num_inference_steps",
-            "generator",
-            "latents",
-            "return_dict",
-            "callback_on_step_end",
-            "callback_on_step_end_tensor_inputs",
-        ]
-    )
-    test_xformers_attention = False
-    supports_dduf = False
-
-    def get_dummy_components(self):
-        torch.manual_seed(0)
-        vae = AutoencoderKLWan(
-            base_dim=3,
-            z_dim=16,
-            dim_mult=[1, 1, 1, 1],
-            num_res_blocks=1,
-            temperal_downsample=[False, True, True],
-        )
-
-        torch.manual_seed(0)
-        scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0)
-        text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5")
-        tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
-
-        torch.manual_seed(0)
-        transformer = WanVACETransformer3DModel(
-            patch_size=(1, 2, 2),
-            num_attention_heads=2,
-            attention_head_dim=12,
-            in_channels=16,
-            out_channels=16,
-            text_dim=32,
-            freq_dim=256,
-            ffn_dim=32,
-            num_layers=3,
-            cross_attn_norm=True,
-            qk_norm="rms_norm_across_heads",
-            rope_max_seq_len=32,
-            vace_layers=[0, 2],
-            vace_in_channels=96,
-        )
-
-        components = {
-            "transformer": transformer,
-            "vae": vae,
-            "scheduler": scheduler,
-            "text_encoder": text_encoder,
-            "tokenizer": tokenizer,
-        }
-        return components
-
-    def get_dummy_inputs(self, device, seed=0):
-        if str(device).startswith("mps"):
-            generator = torch.manual_seed(seed)
-        else:
-            generator = torch.Generator(device=device).manual_seed(seed)
-
-        num_frames = 17
-        height = 16
-        width = 16
-
-        video = [Image.new("RGB", (height, width))] * num_frames
-        mask = [Image.new("L", (height, width), 0)] * num_frames
-
-        inputs = {
-            "video": video,
-            "mask": mask,
-            "prompt": "dance monkey",
-            "negative_prompt": "negative",
-            "generator": generator,
-            "num_inference_steps": 2,
-            "guidance_scale": 6.0,
-            "height": 16,
-            "width": 16,
-            "num_frames": num_frames,
-            "max_sequence_length": 16,
-            "output_type": "pt",
-        }
-        return inputs
-
-    def test_inference(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        video = pipe(**inputs).frames[0]
-        self.assertEqual(video.shape, (17, 3, 16, 16))
-
-        # fmt: off
-        expected_slice = [0.4523, 0.45198, 0.44872, 0.45326, 0.45211, 0.45258, 0.45344, 0.453, 0.52431, 0.52572, 0.50701, 0.5118, 0.53717, 0.53093, 0.50557, 0.51402]
-        # fmt: on
-
-        video_slice = video.flatten()
-        video_slice = torch.cat([video_slice[:8], video_slice[-8:]])
-        video_slice = [round(x, 5) for x in video_slice.tolist()]
-        self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3))
-
-    def test_inference_with_single_reference_image(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        inputs["reference_images"] = Image.new("RGB", (16, 16))
-        video = pipe(**inputs).frames[0]
-        self.assertEqual(video.shape, (17, 3, 16, 16))
-
-        # fmt: off
-        expected_slice = [0.45247, 0.45214, 0.44874, 0.45314, 0.45171, 0.45299, 0.45428, 0.45317, 0.51378, 0.52658, 0.53361, 0.52303, 0.46204, 0.50435, 0.52555, 0.51342]
-        # fmt: on
-
-        video_slice = video.flatten()
-        video_slice = torch.cat([video_slice[:8], video_slice[-8:]])
-        video_slice = [round(x, 5) for x in video_slice.tolist()]
-        self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3))
-
-    def test_inference_with_multiple_reference_image(self):
-        device = "cpu"
-
-        components = self.get_dummy_components()
-        pipe = self.pipeline_class(**components)
-        pipe.to(device)
-        pipe.set_progress_bar_config(disable=None)
-
-        inputs = self.get_dummy_inputs(device)
-        inputs["reference_images"] = [[Image.new("RGB", (16, 16))] * 2]
-        video = pipe(**inputs).frames[0]
-        self.assertEqual(video.shape, (17, 3, 16, 16))
-
-        # fmt: off
-        expected_slice = [0.45321, 0.45221, 0.44818, 0.45375, 0.45268, 0.4519, 0.45271, 0.45253, 0.51244, 0.52223, 0.51253, 0.51321, 0.50743, 0.51177, 0.51626, 0.50983]
-        # fmt: on
-
-        video_slice = video.flatten()
-        video_slice = torch.cat([video_slice[:8], video_slice[-8:]])
-        video_slice = [round(x, 5) for x in video_slice.tolist()]
-        self.assertTrue(np.allclose(video_slice, expected_slice, atol=1e-3))
-
-    @unittest.skip("Test not supported")
-    def test_attention_slicing_forward_pass(self):
-        pass
-
-    @unittest.skip("Errors out because passing multiple prompts at once is not yet supported by this pipeline.")
-    def test_encode_prompt_works_in_isolation(self):
-        pass
-
-    @unittest.skip("Batching is not yet supported with this pipeline")
-    def test_inference_batch_consistent(self):
-        pass
-
-    @unittest.skip("Batching is not yet supported with this pipeline")
-    def test_inference_batch_single_identical(self):
-        return super().test_inference_batch_single_identical()

tests/quantization/bnb/test_4bit.py

@@ -205,7 +205,7 @@ class BnB4BitBasicTests(Base4bitTests):
 
     def test_original_dtype(self):
         r"""
-        A simple test to check if the model successfully stores the original dtype
+        A simple test to check if the model succesfully stores the original dtype
         """
         self.assertTrue("_pre_quantization_dtype" in self.model_4bit.config)
         self.assertFalse("_pre_quantization_dtype" in self.model_fp16.config)
@@ -476,7 +476,6 @@ class SlowBnb4BitTests(Base4bitTests):
         r"""
         Test that loading the model and unquantize it produce correct results.
         """
-        torch.use_deterministic_algorithms(True)
         self.pipeline_4bit.transformer.dequantize()
         output = self.pipeline_4bit(
             prompt=self.prompt,
@@ -527,7 +526,7 @@ class SlowBnb4BitTests(Base4bitTests):
         reason="Test will pass after https://github.com/huggingface/accelerate/pull/3223 is in a release.",
         strict=True,
     )
-    def test_pipeline_cuda_placement_works_with_nf4(self):
+    def test_pipeline_device_placement_works_with_nf4(self):
         transformer_nf4_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -561,7 +560,7 @@ class SlowBnb4BitTests(Base4bitTests):
         ).to(torch_device)
 
         # Check if inference works.
-        _ = pipeline_4bit(self.prompt, max_sequence_length=20, num_inference_steps=2)
+        _ = pipeline_4bit("table", max_sequence_length=20, num_inference_steps=2)
 
         del pipeline_4bit
 

tests/quantization/bnb/test_mixed_int8.py

@@ -195,7 +195,7 @@ class BnB8bitBasicTests(Base8bitTests):
 
     def test_original_dtype(self):
         r"""
-        A simple test to check if the model successfully stores the original dtype
+        A simple test to check if the model succesfully stores the original dtype
         """
         self.assertTrue("_pre_quantization_dtype" in self.model_8bit.config)
         self.assertFalse("_pre_quantization_dtype" in self.model_fp16.config)
@@ -478,7 +478,6 @@ class SlowBnb8bitTests(Base8bitTests):
         r"""
         Test that loading the model and unquantize it produce correct results.
         """
-        torch.use_deterministic_algorithms(True)
         self.pipeline_8bit.transformer.dequantize()
         output = self.pipeline_8bit(
             prompt=self.prompt,
@@ -493,7 +492,7 @@ class SlowBnb8bitTests(Base8bitTests):
         self.assertTrue(max_diff < 1e-2)
 
         # 8bit models cannot be offloaded to CPU.
-        self.assertTrue(self.pipeline_8bit.transformer.device.type == torch_device)
+        self.assertTrue(self.pipeline_8bit.transformer.device.type == "cuda")
         # calling it again shouldn't be a problem
         _ = self.pipeline_8bit(
             prompt=self.prompt,
@@ -535,7 +534,7 @@ class SlowBnb8bitTests(Base8bitTests):
         ).to(device)
 
         # Check if inference works.
-        _ = pipeline_8bit(self.prompt, max_sequence_length=20, num_inference_steps=2)
+        _ = pipeline_8bit("table", max_sequence_length=20, num_inference_steps=2)
 
         del pipeline_8bit
 

tests/single_file/test_model_autoencoder_dc_single_file.py

@@ -95,7 +95,7 @@ class AutoencoderDCSingleFileTests(unittest.TestCase):
         # `in` variant checkpoints require passing in a `config` parameter
         # in order to set the scaling factor correctly.
         # `in` and `mix` variants have the same keys and we cannot automatically infer a scaling factor.
-        # We default to using the `mix` config
+        # We default to using teh `mix` config
         repo_id = "mit-han-lab/dc-ae-f128c512-in-1.0-diffusers"
         ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f128c512-in-1.0/blob/main/model.safetensors"
 

utils/consolidated_test_report.py

@@ -1,789 +0,0 @@
-#!/usr/bin/env python
-import argparse
-import glob
-import os
-import re
-from datetime import date, datetime
-
-from slack_sdk import WebClient
-from tabulate import tabulate
-
-
-MAX_LEN_MESSAGE = 3001  # slack endpoint has a limit of 3001 characters
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--slack_channel_name", default="diffusers-ci-nightly")
-parser.add_argument(
-    "--reports_dir",
-    default="reports",
-    help="Directory containing test reports (will search recursively in all subdirectories)",
-)
-parser.add_argument("--output_file", default=None, help="Path to save the consolidated report (markdown format)")
-
-
-def parse_stats_file(file_path):
-    """Parse a stats file to extract test statistics."""
-    try:
-        with open(file_path, "r") as f:
-            content = f.read()
-
-            # Extract the numbers using regex
-            tests_pattern = r"collected (\d+) items"
-            passed_pattern = r"(\d+) passed"
-            failed_pattern = r"(\d+) failed"
-            skipped_pattern = r"(\d+) skipped"
-            xpassed_pattern = r"(\d+) xpassed"
-
-            tests_match = re.search(tests_pattern, content)
-            passed_match = re.search(passed_pattern, content)
-            failed_match = re.search(failed_pattern, content)
-            skipped_match = re.search(skipped_pattern, content)
-            xpassed_match = re.search(xpassed_pattern, content)
-
-            passed = int(passed_match.group(1)) if passed_match else 0
-            failed = int(failed_match.group(1)) if failed_match else 0
-            skipped = int(skipped_match.group(1)) if skipped_match else 0
-            xpassed = int(xpassed_match.group(1)) if xpassed_match else 0
-
-            # If tests_match exists, use it, otherwise calculate from passed/failed/skipped
-            if tests_match:
-                tests = int(tests_match.group(1))
-            else:
-                tests = passed + failed + skipped + xpassed
-
-            # Extract timing information if available
-            timing_pattern = r"slowest \d+ test durations[\s\S]*?\n([\s\S]*?)={70}"
-            timing_match = re.search(timing_pattern, content, re.MULTILINE)
-            slowest_tests = []
-
-            if timing_match:
-                timing_text = timing_match.group(1).strip()
-                test_timing_lines = timing_text.split("\n")
-                for line in test_timing_lines:
-                    if line.strip():
-                        # Format is typically: 10.37s call     tests/path/to/test.py::TestClass::test_method
-                        parts = line.strip().split()
-                        if len(parts) >= 3:
-                            time_str = parts[0]
-                            test_path = " ".join(parts[2:])
-
-                            # Skip entries with "< 0.05 secs were omitted" or similar
-                            if "secs were omitted" in test_path:
-                                continue
-
-                            try:
-                                time_seconds = float(time_str.rstrip("s"))
-                                slowest_tests.append({"test": test_path, "duration": time_seconds})
-                            except ValueError:
-                                pass
-
-            return {
-                "tests": tests,
-                "passed": passed,
-                "failed": failed,
-                "skipped": skipped,
-                "slowest_tests": slowest_tests,
-            }
-    except Exception as e:
-        print(f"Error parsing {file_path}: {e}")
-        return {"tests": 0, "passed": 0, "failed": 0, "skipped": 0, "slowest_tests": []}
-
-
-def parse_durations_file(file_path):
-    """Parse a durations file to extract test timing information."""
-    slowest_tests = []
-    try:
-        durations_file = file_path.replace("_stats.txt", "_durations.txt")
-        if os.path.exists(durations_file):
-            with open(durations_file, "r") as f:
-                content = f.read()
-
-                # Skip the header line
-                for line in content.split("\n")[1:]:
-                    if line.strip():
-                        # Format is typically: 10.37s call     tests/path/to/test.py::TestClass::test_method
-                        parts = line.strip().split()
-                        if len(parts) >= 3:
-                            time_str = parts[0]
-                            test_path = " ".join(parts[2:])
-
-                            # Skip entries with "< 0.05 secs were omitted" or similar
-                            if "secs were omitted" in test_path:
-                                continue
-
-                            try:
-                                time_seconds = float(time_str.rstrip("s"))
-                                slowest_tests.append({"test": test_path, "duration": time_seconds})
-                            except ValueError:
-                                # If time_str is not a valid float, it might be a different format
-                                # For example, some pytest formats show "< 0.05s" or similar
-                                if test_path.startswith("<") and "secs were omitted" in test_path:
-                                    # Extract the time value from test_path if it's in the format "< 0.05 secs were omitted"
-                                    try:
-                                        # This handles entries where the time is in the test_path itself
-                                        dur_match = re.search(r"(\d+(?:\.\d+)?)", test_path)
-                                        if dur_match:
-                                            time_seconds = float(dur_match.group(1))
-                                            slowest_tests.append({"test": test_path, "duration": time_seconds})
-                                    except ValueError:
-                                        pass
-    except Exception as e:
-        print(f"Error parsing durations file {file_path.replace('_stats.txt', '_durations.txt')}: {e}")
-
-    return slowest_tests
-
-
-def parse_failures_file(file_path):
-    """Parse a failures file to extract failed test details."""
-    failures = []
-    try:
-        with open(file_path, "r") as f:
-            content = f.read()
-
-            # We don't need the base file name anymore as we're getting test paths from summary
-
-            # Check if it's a short stack format
-            if "============================= FAILURES SHORT STACK =============================" in content:
-                # First, look for pytest-style failure headers with underscores and clean them up
-                test_headers = re.findall(r"_{5,}\s+([^_\n]+?)\s+_{5,}", content)
-
-                for test_name in test_headers:
-                    test_name = test_name.strip()
-                    # Make sure it's a valid test name (contains a dot and doesn't look like a number)
-                    if "." in test_name and not test_name.replace(".", "").isdigit():
-                        # For test names missing the full path, check if we can reconstruct it from failures_line.txt
-                        # This is a best effort - we won't always have the line file available
-                        if not test_name.endswith(".py") and "::" not in test_name and "/" not in test_name:
-                            # Try to look for a corresponding line file
-                            line_file = file_path.replace("_failures_short.txt", "_failures_line.txt")
-                            if os.path.exists(line_file):
-                                try:
-                                    with open(line_file, "r") as lf:
-                                        line_content = lf.read()
-                                        # Look for test name in line file which might have the full path
-                                        path_match = re.search(
-                                            r"(tests/[\w/]+\.py::[^:]+::" + test_name.split(".")[-1] + ")",
-                                            line_content,
-                                        )
-                                        if path_match:
-                                            test_name = path_match.group(1)
-                                except Exception:
-                                    pass  # If we can't read the line file, just use what we have
-
-                        failures.append(
-                            {
-                                "test": test_name,
-                                "error": "Error occurred",
-                                "original_test_name": test_name,  # Keep original for reference
-                            }
-                        )
-
-                # If we didn't find any pytest-style headers, try other formats
-                if not failures:
-                    # Look for test names at the beginning of the file (in first few lines)
-                    first_lines = content.split("\n")[:20]  # Look at first 20 lines
-                    for line in first_lines:
-                        # Look for test names in various formats
-                        # Format: tests/file.py::TestClass::test_method
-                        path_match = re.search(r"(tests/[\w/]+\.py::[\w\.]+::\w+)", line)
-                        # Format: TestClass.test_method
-                        class_match = re.search(r"([A-Za-z][A-Za-z0-9_]+\.[A-Za-z][A-Za-z0-9_]+)", line)
-
-                        if path_match:
-                            test_name = path_match.group(1)
-                            failures.append(
-                                {"test": test_name, "error": "Error occurred", "original_test_name": test_name}
-                            )
-                            break  # Found a full path, stop looking
-                        elif class_match and "test" in line.lower():
-                            test_name = class_match.group(1)
-                            # Make sure it's likely a test name (contains test in method name)
-                            if "test" in test_name.lower():
-                                failures.append(
-                                    {"test": test_name, "error": "Error occurred", "original_test_name": test_name}
-                                )
-            else:
-                # Standard format - try to extract from standard pytest output
-                failure_blocks = re.split(r"={70}", content)
-
-                for block in failure_blocks:
-                    if not block.strip():
-                        continue
-
-                    # Look for test paths in the format: path/to/test.py::TestClass::test_method
-                    path_matches = re.findall(r"([\w/]+\.py::[\w\.]+::\w+)", block)
-                    if path_matches:
-                        for test_name in path_matches:
-                            failures.append(
-                                {"test": test_name, "error": "Error occurred", "original_test_name": test_name}
-                            )
-                    else:
-                        # Try alternative format: TestClass.test_method
-                        class_matches = re.findall(r"([A-Za-z][A-Za-z0-9_]+\.[A-Za-z][A-Za-z0-9_]+)", block)
-                        for test_name in class_matches:
-                            # Filter out things that don't look like test names
-                            if (
-                                not test_name.startswith(("e.g", "i.e", "etc."))
-                                and not test_name.isdigit()
-                                and "test" in test_name.lower()
-                            ):
-                                failures.append(
-                                    {"test": test_name, "error": "Error occurred", "original_test_name": test_name}
-                                )
-
-    except Exception as e:
-        print(f"Error parsing failures in {file_path}: {e}")
-
-    return failures
-
-
-def consolidate_reports(reports_dir):
-    """Consolidate test reports from multiple test runs, including from subdirectories."""
-    # Get all stats files, including those in subdirectories
-    stats_files = glob.glob(f"{reports_dir}/**/*_stats.txt", recursive=True)
-
-    results = {}
-    total_stats = {"tests": 0, "passed": 0, "failed": 0, "skipped": 0}
-
-    # Collect all slow tests across all test suites
-    all_slow_tests = []
-
-    # Process each stats file and its corresponding failures file
-    for stats_file in stats_files:
-        # Extract test suite name from filename (e.g., tests_pipeline_allegro_cuda_stats.txt -> pipeline_allegro_cuda)
-        base_name = os.path.basename(stats_file).replace("_stats.txt", "")
-
-        # Include parent directory in suite name if it's in a subdirectory
-        rel_path = os.path.relpath(os.path.dirname(stats_file), reports_dir)
-        if rel_path and rel_path != ".":
-            # Remove 'test_reports' suffix from directory name if present
-            dir_name = os.path.basename(rel_path)
-            if dir_name.endswith("_test_reports"):
-                dir_name = dir_name[:-13]  # Remove '_test_reports' suffix
-            base_name = f"{dir_name}/{base_name}"
-
-        # Parse stats
-        stats = parse_stats_file(stats_file)
-
-        # If no slowest tests found in stats file, try the durations file directly
-        if not stats.get("slowest_tests"):
-            stats["slowest_tests"] = parse_durations_file(stats_file)
-
-        # Update total stats
-        for key in ["tests", "passed", "failed", "skipped"]:
-            total_stats[key] += stats[key]
-
-        # Collect slowest tests with their suite name
-        for slow_test in stats.get("slowest_tests", []):
-            all_slow_tests.append({"test": slow_test["test"], "duration": slow_test["duration"], "suite": base_name})
-
-        # Parse failures if there are any
-        failures = []
-        if stats["failed"] > 0:
-            # First try to get test paths from summary_short.txt which has the best format
-            summary_file = stats_file.replace("_stats.txt", "_summary_short.txt")
-            if os.path.exists(summary_file):
-                try:
-                    with open(summary_file, "r") as f:
-                        content = f.read()
-                        # Look for full lines with test path and error message: "FAILED test_path - error_msg"
-                        failed_test_lines = re.findall(
-                            r"FAILED\s+(tests/[\w/]+\.py::[A-Za-z0-9_\.]+::[A-Za-z0-9_]+)(?:\s+-\s+(.+))?", content
-                        )
-
-                        if failed_test_lines:
-                            for match in failed_test_lines:
-                                test_path = match[0]
-                                error_msg = match[1] if len(match) > 1 and match[1] else "No error message"
-
-                                failures.append({"test": test_path, "error": error_msg})
-                except Exception as e:
-                    print(f"Error parsing summary file: {e}")
-
-            # If no failures found in summary, try other failure files
-            if not failures:
-                failure_patterns = ["_failures_short.txt", "_failures.txt", "_failures_line.txt", "_failures_long.txt"]
-
-                for pattern in failure_patterns:
-                    failures_file = stats_file.replace("_stats.txt", pattern)
-                    if os.path.exists(failures_file):
-                        failures = parse_failures_file(failures_file)
-                        if failures:
-                            break
-
-                # No debug output needed
-
-        # Store results for this test suite
-        results[base_name] = {"stats": stats, "failures": failures}
-
-    # Filter out entries with "secs were omitted"
-    filtered_slow_tests = [test for test in all_slow_tests if "secs were omitted" not in test["test"]]
-
-    # Sort all slow tests by duration (descending)
-    filtered_slow_tests.sort(key=lambda x: x["duration"], reverse=True)
-
-    # Get the number of slowest tests to show from environment variable or default to 10
-    num_slowest_tests = int(os.environ.get("SHOW_SLOWEST_TESTS", "10"))
-    top_slowest_tests = filtered_slow_tests[:num_slowest_tests] if filtered_slow_tests else []
-
-    # Calculate additional duration statistics
-    total_duration = sum(test["duration"] for test in all_slow_tests)
-
-    # Calculate duration per suite
-    suite_durations = {}
-    for test in all_slow_tests:
-        suite_name = test["suite"]
-        if suite_name not in suite_durations:
-            suite_durations[suite_name] = 0
-        suite_durations[suite_name] += test["duration"]
-
-    # Removed duration categories
-
-    return {
-        "total_stats": total_stats,
-        "test_suites": results,
-        "slowest_tests": top_slowest_tests,
-        "duration_stats": {"total_duration": total_duration, "suite_durations": suite_durations},
-    }
-
-
-def generate_report(consolidated_data):
-    """Generate a comprehensive markdown report from consolidated data."""
-    report = []
-
-    # Add report header
-    report.append("# Diffusers Nightly Test Report")
-    report.append(f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
-
-    # Removed comparison section
-
-    # Add summary section
-    total = consolidated_data["total_stats"]
-    report.append("## Summary")
-
-    # Get duration stats if available
-    duration_stats = consolidated_data.get("duration_stats", {})
-    total_duration = duration_stats.get("total_duration", 0)
-
-    summary_table = [
-        ["Total Tests", total["tests"]],
-        ["Passed", total["passed"]],
-        ["Failed", total["failed"]],
-        ["Skipped", total["skipped"]],
-        ["Success Rate", f"{(total['passed'] / total['tests'] * 100):.2f}%" if total["tests"] > 0 else "N/A"],
-        ["Total Duration", f"{total_duration:.2f}s" if total_duration else "N/A"],
-    ]
-
-    report.append(tabulate(summary_table, tablefmt="pipe"))
-    report.append("")
-
-    # Removed duration distribution section
-
-    # Add test suites summary
-    report.append("## Test Suites")
-
-    # Include duration in test suites table if available
-    suite_durations = consolidated_data.get("duration_stats", {}).get("suite_durations", {})
-
-    if suite_durations:
-        suites_table = [["Test Suite", "Tests", "Passed", "Failed", "Skipped", "Success Rate", "Duration (s)"]]
-    else:
-        suites_table = [["Test Suite", "Tests", "Passed", "Failed", "Skipped", "Success Rate"]]
-
-    # Sort test suites by success rate (ascending - least successful first)
-    sorted_suites = sorted(
-        consolidated_data["test_suites"].items(),
-        key=lambda x: (x[1]["stats"]["passed"] / x[1]["stats"]["tests"] * 100) if x[1]["stats"]["tests"] > 0 else 0,
-        reverse=False,
-    )
-
-    for suite_name, suite_data in sorted_suites:
-        stats = suite_data["stats"]
-        success_rate = f"{(stats['passed'] / stats['tests'] * 100):.2f}%" if stats["tests"] > 0 else "N/A"
-
-        if suite_durations:
-            duration = suite_durations.get(suite_name, 0)
-            suites_table.append(
-                [
-                    suite_name,
-                    stats["tests"],
-                    stats["passed"],
-                    stats["failed"],
-                    stats["skipped"],
-                    success_rate,
-                    f"{duration:.2f}",
-                ]
-            )
-        else:
-            suites_table.append(
-                [suite_name, stats["tests"], stats["passed"], stats["failed"], stats["skipped"], success_rate]
-            )
-
-    report.append(tabulate(suites_table, headers="firstrow", tablefmt="pipe"))
-    report.append("")
-
-    # Add slowest tests section
-    slowest_tests = consolidated_data.get("slowest_tests", [])
-    if slowest_tests:
-        report.append("## Slowest Tests")
-
-        slowest_table = [["Rank", "Test", "Duration (s)", "Test Suite"]]
-        for i, test in enumerate(slowest_tests, 1):
-            # Skip entries that don't contain actual test names
-            if "< 0.05 secs were omitted" in test["test"]:
-                continue
-            slowest_table.append([i, test["test"], f"{test['duration']:.2f}", test["suite"]])
-
-        report.append(tabulate(slowest_table, headers="firstrow", tablefmt="pipe"))
-        report.append("")
-
-    # Add failures section if there are any
-    failed_suites = [s for s in sorted_suites if s[1]["stats"]["failed"] > 0]
-
-    if failed_suites:
-        report.append("## Failures")
-
-        # Group failures by module for cleaner organization
-        failures_by_module = {}
-
-        for suite_name, suite_data in failed_suites:
-            # Extract failures data for this suite
-            for failure in suite_data.get("failures", []):
-                test_name = failure["test"]
-
-                # If test name doesn't look like a full path, try to reconstruct it
-                if not ("/" in test_name or "::" in test_name) and "." in test_name:
-                    # For simple 'TestClass.test_method' format, try to get full path from suite name
-                    # Form: tests_<suite>_cuda -> tests/<suite>/test_<suite>.py::TestClass::test_method
-                    if suite_name.startswith("tests_") and "_cuda" in suite_name:
-                        # Extract component name from suite
-                        component = suite_name.replace("tests_", "").replace("_cuda", "")
-                        if "." in test_name:
-                            class_name, method_name = test_name.split(".", 1)
-                            possible_path = f"tests/{component}/test_{component}.py::{class_name}::{method_name}"
-                            # Use this constructed path if it seems reasonable
-                            if "test_" in method_name:
-                                test_name = possible_path
-
-                # Extract module name from test name
-                if "::" in test_name:
-                    # For path/file.py::TestClass::test_method format
-                    parts = test_name.split("::")
-                    module_name = parts[-2] if len(parts) >= 2 else "Other"  # TestClass
-                elif "." in test_name:
-                    # For TestClass.test_method format
-                    parts = test_name.split(".")
-                    module_name = parts[0]  # TestClass
-                else:
-                    module_name = "Other"
-
-                # Skip module names that don't look like class/module names
-                if (
-                    module_name.startswith(("e.g", "i.e", "etc"))
-                    or module_name.replace(".", "").isdigit()
-                    or len(module_name) < 3
-                ):
-                    module_name = "Other"
-
-                # Add to the module group
-                if module_name not in failures_by_module:
-                    failures_by_module[module_name] = []
-
-                # Prepend the suite name if the test name doesn't already have a full path
-                if "/" not in test_name and suite_name not in test_name:
-                    full_test_name = f"{suite_name}::{test_name}"
-                else:
-                    full_test_name = test_name
-
-                # Add this failure to the module group
-                failures_by_module[module_name].append(
-                    {"test": full_test_name, "original_test": test_name, "error": failure["error"]}
-                )
-
-        # Create a list of failing tests for each module
-        if failures_by_module:
-            for module_name, failures in sorted(failures_by_module.items()):
-                report.append(f"### {module_name}")
-
-                # Put all failed tests in a single code block
-                report.append("```")
-                for failure in failures:
-                    # Show test path and error message if available
-                    if failure.get("error") and failure["error"] != "No error message":
-                        report.append(f"{failure['test']} - {failure['error']}")
-                    else:
-                        report.append(failure["test"])
-                report.append("```")
-
-                report.append("")  # Add space between modules
-        else:
-            report.append("*No detailed failure information available*")
-            report.append("")
-
-    return "\n".join(report)
-
-
-def create_test_groups_table(test_groups, total_tests, total_success_rate):
-    """Create a table-like format for test groups showing total tests and success rate."""
-    if not test_groups:
-        return None
-
-    # Sort by total test count (descending)
-    sorted_groups = sorted(test_groups.items(), key=lambda x: x[1]["total"], reverse=True)
-
-    # Create table lines
-    table_lines = ["```"]
-    table_lines.append("Test Results Summary")
-    table_lines.append("-------------------")
-    table_lines.append(f"Total Tests:  {total_tests:,}")
-    table_lines.append(f"Success Rate: {total_success_rate}")
-    table_lines.append("")
-    table_lines.append("Category            | Total Tests | Failed | Success Rate")
-    table_lines.append("------------------- | ----------- | ------ | ------------")
-
-    # Add rows
-    for category, stats in sorted_groups:
-        # Pad category name to fixed width (19 chars)
-        padded_cat = category[:19].ljust(19)  # Truncate if too long
-        # Right-align counts
-        padded_total = str(stats["total"]).rjust(11)
-        padded_failed = str(stats["failed"]).rjust(6)
-        # Calculate and format success rate
-        if stats["total"] > 0:
-            cat_success_rate = f"{((stats['total'] - stats['failed']) / stats['total'] * 100):.1f}%"
-        else:
-            cat_success_rate = "N/A"
-        padded_rate = cat_success_rate.rjust(12)
-        table_lines.append(f"{padded_cat} | {padded_total} | {padded_failed} | {padded_rate}")
-
-    table_lines.append("```")
-
-    total_failures = sum(stats["failed"] for stats in test_groups.values())
-    return (
-        f"*Test Groups Summary ({total_failures} {'failure' if total_failures == 1 else 'failures'}):*\n"
-        + "\n".join(table_lines)
-    )
-
-
-def create_slack_payload(consolidated_data):
-    """Create a concise Slack message payload from consolidated data."""
-    total = consolidated_data["total_stats"]
-    success_rate = f"{(total['passed'] / total['tests'] * 100):.2f}%" if total["tests"] > 0 else "N/A"
-
-    # Determine emoji based on success rate
-    if total["failed"] == 0:
-        emoji = "✅"
-    elif total["failed"] / total["tests"] < 0.1:
-        emoji = "⚠️"
-    else:
-        emoji = "❌"
-
-    # Create a more compact summary section
-    summary = f"{emoji} *Diffusers Nightly Tests:* {success_rate} success ({total['passed']}/{total['tests']} tests"
-    if total["skipped"] > 0:
-        summary += f", {total['skipped']} skipped"
-    summary += ")"
-
-    # Create the test suites table in markdown format
-    # Build the markdown table with proper alignment
-    table_lines = []
-    table_lines.append("```")
-
-    # Sort test suites by success rate (ascending - least successful first)
-    sorted_suites = sorted(
-        consolidated_data["test_suites"].items(),
-        key=lambda x: (x[1]["stats"]["passed"] / x[1]["stats"]["tests"] * 100) if x[1]["stats"]["tests"] > 0 else 0,
-        reverse=False,
-    )
-
-    # Calculate max widths for proper alignment
-    max_suite_name_len = max(len(suite_name) for suite_name, _ in sorted_suites) if sorted_suites else 10
-    max_suite_name_len = max(max_suite_name_len, len("Test Suite"))  # Ensure header fits
-
-    # Create header with proper spacing (only Tests, Failed, Success Rate)
-    header = f"| {'Test Suite'.ljust(max_suite_name_len)} | {'Tests'.rjust(6)} | {'Failed'.rjust(6)} | {'Success Rate'.ljust(12)} |"
-    separator = f"|:{'-' * max_suite_name_len}|{'-' * 7}:|{'-' * 7}:|:{'-' * 11}|"
-
-    table_lines.append(header)
-    table_lines.append(separator)
-
-    # Add data rows with proper alignment
-    for suite_name, suite_data in sorted_suites:
-        stats = suite_data["stats"]
-        suite_success_rate = f"{(stats['passed'] / stats['tests'] * 100):.2f}%" if stats["tests"] > 0 else "N/A"
-
-        row = f"| {suite_name.ljust(max_suite_name_len)} | {str(stats['tests']).rjust(6)} | {str(stats['failed']).rjust(6)} | {suite_success_rate.ljust(12)} |"
-
-        table_lines.append(row)
-
-    table_lines.append("```")
-
-    # Create the Slack payload with character limit enforcement
-    payload = [
-        {"type": "section", "text": {"type": "mrkdwn", "text": summary}},
-        {"type": "section", "text": {"type": "mrkdwn", "text": "\n".join(table_lines)}},
-    ]
-
-    # Add action button
-    if os.environ.get("GITHUB_RUN_ID"):
-        run_id = os.environ["GITHUB_RUN_ID"]
-        payload.append(
-            {
-                "type": "section",
-                "text": {
-                    "type": "mrkdwn",
-                    "text": f"*<https://github.com/huggingface/diffusers/actions/runs/{run_id}|View full report on GitHub>*",
-                },
-            }
-        )
-
-    # Add date in more compact form
-    payload.append(
-        {
-            "type": "context",
-            "elements": [
-                {
-                    "type": "plain_text",
-                    "text": f"Results for {date.today()}",
-                },
-            ],
-        }
-    )
-
-    # Enforce 3001 character limit
-    payload_text = str(payload)
-    if len(payload_text) > MAX_LEN_MESSAGE:
-        # Truncate table if payload is too long
-        # Remove rows from the bottom until under limit
-        original_table_lines = table_lines[:]
-        while len(str(payload)) > MAX_LEN_MESSAGE and len(table_lines) > 3:  # Keep at least header and separator
-            # Remove the last data row (but keep ``` at the end)
-            table_lines.pop(-2)  # Remove second to last (last is the closing ```)
-
-            # Recreate payload with truncated table
-            payload[1] = {"type": "section", "text": {"type": "mrkdwn", "text": "\n".join(table_lines)}}
-
-        # Add note if we had to truncate
-        if len(table_lines) < len(original_table_lines):
-            truncated_count = len(original_table_lines) - len(table_lines)
-            table_lines.insert(-1, f"... {truncated_count} more test suites (truncated due to message limit)")
-            payload[1] = {"type": "section", "text": {"type": "mrkdwn", "text": "\n".join(table_lines)}}
-
-    return payload
-
-
-def create_failed_tests_by_suite_ordered(consolidated_data):
-    """Group failed tests by test suite, ordered by success rate (ascending)."""
-    # Sort test suites by success rate (ascending - least successful first)
-    sorted_suites = sorted(
-        consolidated_data["test_suites"].items(),
-        key=lambda x: (x[1]["stats"]["passed"] / x[1]["stats"]["tests"] * 100) if x[1]["stats"]["tests"] > 0 else 0,
-        reverse=False,
-    )
-
-    failed_suite_tests = []
-
-    # Process suites in order of success rate
-    for suite_name, suite_data in sorted_suites:
-        if suite_data["stats"]["failed"] > 0:
-            suite_failures = []
-
-            for failure in suite_data.get("failures", []):
-                test_name = failure["test"]
-
-                # Try to reconstruct full path if partial
-                if "::" in test_name and "/" in test_name:
-                    full_test_name = test_name
-                elif "::" in test_name or "." in test_name:
-                    if "/" not in test_name and suite_name not in test_name:
-                        full_test_name = f"{suite_name}::{test_name}"
-                    else:
-                        full_test_name = test_name
-                else:
-                    full_test_name = f"{suite_name}::{test_name}"
-
-                suite_failures.append(full_test_name)
-
-            # Sort and deduplicate tests within the suite
-            suite_failures = sorted(set(suite_failures))
-
-            if suite_failures:
-                failed_suite_tests.append(
-                    {
-                        "suite_name": suite_name,
-                        "tests": suite_failures,
-                        "success_rate": (suite_data["stats"]["passed"] / suite_data["stats"]["tests"] * 100)
-                        if suite_data["stats"]["tests"] > 0
-                        else 0,
-                    }
-                )
-
-    return failed_suite_tests
-
-
-def main(args):
-    # Make sure reports directory exists
-    if not os.path.isdir(args.reports_dir):
-        print(f"Error: Reports directory '{args.reports_dir}' does not exist.")
-        return
-
-    # Consolidate reports
-    consolidated_data = consolidate_reports(args.reports_dir)
-
-    # Check if we found any test results
-    if consolidated_data["total_stats"]["tests"] == 0:
-        print(f"Warning: No test results found in '{args.reports_dir}' or its subdirectories.")
-
-    # Generate markdown report
-    report = generate_report(consolidated_data)
-
-    # Save report to file if specified
-    if args.output_file:
-        # Create parent directories if they don't exist
-        output_dir = os.path.dirname(args.output_file)
-        if output_dir and not os.path.exists(output_dir):
-            os.makedirs(output_dir)
-
-        with open(args.output_file, "w") as f:
-            f.write(report)
-
-        # Only print the report when saving to file
-        print(report)
-
-    # Send to Slack if token is available (optional, can be disabled)
-    slack_token = os.environ.get("SLACK_API_TOKEN")
-    if slack_token and args.slack_channel_name:
-        payload = create_slack_payload(consolidated_data)
-
-        try:
-            client = WebClient(token=slack_token)
-            # Send main message
-            response = client.chat_postMessage(channel=f"#{args.slack_channel_name}", blocks=payload)
-            print(f"Report sent to Slack channel: {args.slack_channel_name}")
-
-            # Send failed tests as separate threaded replies grouped by test suite (ordered by success rate)
-            total = consolidated_data["total_stats"]
-            if total["failed"] > 0:
-                failed_suites = create_failed_tests_by_suite_ordered(consolidated_data)
-                for suite_info in failed_suites:
-                    suite_name = suite_info["suite_name"]
-                    suite_tests = suite_info["tests"]
-                    success_rate = suite_info["success_rate"]
-                    message_text = (
-                        f"**{suite_name}** (Success Rate: {success_rate:.2f}%)\n```\n"
-                        + "\n".join(suite_tests)
-                        + "\n```"
-                    )
-                    client.chat_postMessage(
-                        channel=f"#{args.slack_channel_name}",
-                        thread_ts=response["ts"],  # Reply in thread
-                        text=message_text,  # Use text instead of blocks for markdown
-                    )
-                print(f"Failed tests details sent as {len(failed_suites)} thread replies")
-        except Exception as e:
-            print(f"Error sending report to Slack: {e}")
-
-
-if __name__ == "__main__":
-    args = parser.parse_args()
-    main(args)

utils/custom_init_isort.py

@@ -252,7 +252,7 @@ def sort_imports(file: str, check_only: bool = True):
         code, start_prompt="_import_structure = {", end_prompt="if TYPE_CHECKING:"
     )
 
-    # We ignore block 0 (everything until start_prompt) and the last block (everything after end_prompt).
+    # We ignore block 0 (everything untils start_prompt) and the last block (everything after end_prompt).
     for block_idx in range(1, len(main_blocks) - 1):
         # Check if the block contains some `_import_structure`s thingy to sort.
         block = main_blocks[block_idx]