Merge branch 'main' into controlnet-test-fixes

update
2024-03-15 12:30:25 +05:30 · 2024-03-14 06:53:52 +00:00 · 2024-03-14 06:47:29 +00:00
314 changed files with 3114 additions and 13718 deletions
@@ -12,96 +12,110 @@ env:
  PYTEST_TIMEOUT: 600
  RUN_SLOW: yes
  RUN_NIGHTLY: yes
  PIPELINE_USAGE_CUTOFF: 5000
  SLACK_API_TOKEN: ${{ secrets.SLACK_CIFEEDBACK_BOT_TOKEN }}
 jobs:
-  setup_torch_cuda_pipeline_matrix:
+  run_nightly_tests:
    name: Setup Torch Pipelines Matrix
    runs-on: ubuntu-latest
    outputs:
      pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
    steps:
      - name: Checkout diffusers
        uses: actions/checkout@v3
        with:
          fetch-depth: 2
      - name: Set up Python
        uses: actions/setup-python@v4
        with:
          python-version: "3.8"
      - name: Install dependencies
        run: |
          pip install -e .
          pip install huggingface_hub
      - name: Fetch Pipeline Matrix
        id: fetch_pipeline_matrix
        run: |
          matrix=$(python utils/fetch_torch_cuda_pipeline_test_matrix.py)
          echo $matrix
          echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT
      - name: Pipeline Tests Artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v2
        with:
          name: test-pipelines.json
          path: reports
  run_nightly_tests_for_torch_pipelines:
    name: Torch Pipelines CUDA Nightly Tests
    needs: setup_torch_cuda_pipeline_matrix
    strategy:
      fail-fast: false
      matrix:
-        module: ${{ fromJson(needs.setup_torch_cuda_pipeline_matrix.outputs.pipeline_test_matrix) }}
+        config:
-    runs-on: [single-gpu, nvidia-gpu, t4, ci]
+          - name: Nightly PyTorch CUDA tests on Ubuntu
            framework: pytorch
            runner: docker-gpu
            image: diffusers/diffusers-pytorch-cuda
            report: torch_cuda
          - name: Nightly Flax TPU tests on Ubuntu
            framework: flax
            runner: docker-tpu
            image: diffusers/diffusers-flax-tpu
            report: flax_tpu
          - name: Nightly ONNXRuntime CUDA tests on Ubuntu
            framework: onnxruntime
            runner: docker-gpu
            image: diffusers/diffusers-onnxruntime-cuda
            report: onnx_cuda
    name: ${{ matrix.config.name }}
    runs-on: ${{ matrix.config.runner }}
    container:
-      image: diffusers/diffusers-pytorch-cuda
+      image: ${{ matrix.config.image }}
-      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ --gpus 0
+      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ ${{ matrix.config.runner == 'docker-tpu' && '--privileged' || '--gpus 0'}}
    defaults:
      run:
        shell: bash
    steps:
      - name: Checkout diffusers
        uses: actions/checkout@v3
        with:
          fetch-depth: 2
      - name: NVIDIA-SMI
-        run: nvidia-smi
+        if: ${{ matrix.config.runner == 'docker-gpu' }}
        run: |
          nvidia-smi
      - name: Install dependencies
        run: |
          apt-get update && apt-get install libsndfile1-dev libgl1 -y
          python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
          python -m uv pip install -e [quality,test]
-          python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate.git
+          python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers
          python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate
          python -m uv pip install pytest-reportlog
      - name: Environment
        run: |
          python utils/print_env.py
-      - name: Nightly PyTorch CUDA checkpoint (pipelines) tests 
+      - name: Run nightly PyTorch CUDA tests
        if: ${{ matrix.config.framework == 'pytorch' }}
        env:
          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
          # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
          CUBLAS_WORKSPACE_CONFIG: :16:8
        run: |
          python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
          python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
            -s -v -k "not Flax and not Onnx" \
-            --make-reports=tests_pipeline_${{ matrix.module }}_cuda \
+            --make-reports=tests_${{ matrix.config.report }} \
-            --report-log=tests_pipeline_${{ matrix.module }}_cuda.log \ 
+            --report-log=${{ matrix.config.report }}.log \
-            tests/pipelines/${{ matrix.module }}
+            tests/ 
      - name: Run nightly Flax TPU tests
        if: ${{ matrix.config.framework == 'flax' }}
        env:
          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
        run: |
          python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
          python -m pytest -n 0 \
            -s -v -k "Flax" \
            --make-reports=tests_${{ matrix.config.report }} \
            --report-log=${{ matrix.config.report }}.log \
            tests/
      - name: Run nightly ONNXRuntime CUDA tests
        if: ${{ matrix.config.framework == 'onnxruntime' }}
        env:
          HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
        run: |
          python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
          python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
            -s -v -k "Onnx" \
            --make-reports=tests_${{ matrix.config.report }} \
            --report-log=${{ matrix.config.report }}.log \ 
            tests/
      - name: Failure short reports
        if: ${{ failure() }}
-        run: |
+        run: cat reports/tests_${{ matrix.config.report }}_failures_short.txt
          cat reports/tests_pipeline_${{ matrix.module }}_cuda_stats.txt
          cat reports/tests_pipeline_${{ matrix.module }}_cuda_failures_short.txt
      - name: Test suite reports artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v2
        with:
-          name: pipeline_${{ matrix.module }}_test_reports
+          name: ${{ matrix.config.report }}_test_reports
          path: reports
      - name: Generate Report and Notify Channel
@@ -110,248 +124,6 @@ jobs:
          pip install slack_sdk tabulate
          python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY
  run_nightly_tests_for_other_torch_modules:
    name: Torch Non-Pipelines CUDA Nightly Tests
    runs-on: docker-gpu
    container:
      image: diffusers/diffusers-pytorch-cuda
      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ --gpus 0
    defaults:
      run:
        shell: bash
    strategy:
      matrix:
        module: [models, schedulers, others, examples]
    steps:
    - name: Checkout diffusers
      uses: actions/checkout@v3
      with:
        fetch-depth: 2
    - name: Install dependencies
      run: |
        apt-get update && apt-get install libsndfile1-dev libgl1 -y
        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
        python -m uv pip install -e [quality,test]
        python -m uv pip install 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 PyTorch CUDA tests for non-pipeline modules
      if: ${{ matrix.module != 'examples'}} 
      env:
        HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
        # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
        CUBLAS_WORKSPACE_CONFIG: :16:8
      run: |
        python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
          -s -v -k "not Flax and not Onnx" \
          --make-reports=tests_torch_${{ matrix.module }}_cuda \
          --report-log=tests_torch_${{ matrix.module }}_cuda.log \ 
          tests/${{ matrix.module }}
    - name: Run nightly example tests with Torch
      if: ${{ matrix.module == 'examples' }}
      env:
        HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
        # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
        CUBLAS_WORKSPACE_CONFIG: :16:8
      run: |
        python -m uv pip install peft@git+https://github.com/huggingface/peft.git
        python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
          -s -v --make-reports=examples_torch_cuda \
          --report-log=examples_torch_cuda.log \ 
          examples/
    - name: Failure short reports
      if: ${{ failure() }}
      run: |
        cat reports/tests_torch_${{ matrix.module }}_cuda_stats.txt 
        cat reports/tests_torch_${{ matrix.module }}_cuda_failures_short.txt
    - name: Test suite reports artifacts
      if: ${{ always() }}
      uses: actions/upload-artifact@v2
      with:
        name: torch_${{ matrix.module }}_cuda_test_reports
        path: reports
    - name: Generate Report and Notify Channel
      if: always()
      run: |
        pip install slack_sdk tabulate
        python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY
  run_lora_nightly_tests:
    name: Nightly LoRA Tests with PEFT and TORCH
    runs-on: docker-gpu
    container:
      image: diffusers/diffusers-pytorch-cuda
      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ --gpus 0
    defaults:
      run:
        shell: bash
    steps:
    - name: Checkout diffusers
      uses: actions/checkout@v3
      with:
        fetch-depth: 2
    - name: Install dependencies
      run: |
        apt-get update && apt-get install libsndfile1-dev libgl1 -y
        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
        python -m uv pip install -e [quality,test]
        python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate.git
        python -m uv pip install peft@git+https://github.com/huggingface/peft.git
        python -m uv pip install pytest-reportlog
    - name: Environment
      run: python utils/print_env.py
    - name: Run nightly LoRA tests with PEFT and Torch
      env:
        HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
        # https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
        CUBLAS_WORKSPACE_CONFIG: :16:8
      run: |
        python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
          -s -v -k "not Flax and not Onnx" \
          --make-reports=tests_torch_lora_cuda \
          --report-log=tests_torch_lora_cuda.log \ 
          tests/lora
    - name: Failure short reports
      if: ${{ failure() }}
      run: |
        cat reports/tests_torch_lora_cuda_stats.txt 
        cat reports/tests_torch_lora_cuda_failures_short.txt
    - name: Test suite reports artifacts
      if: ${{ always() }}
      uses: actions/upload-artifact@v2
      with:
        name: torch_lora_cuda_test_reports
        path: reports
    - name: Generate Report and Notify Channel
      if: always()
      run: |
        pip install slack_sdk tabulate
        python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY
  run_flax_tpu_tests:
    name: Nightly Flax TPU Tests
    runs-on: docker-tpu
    container:
      image: diffusers/diffusers-flax-tpu
      options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ --privileged
    defaults:
      run:
        shell: bash
    steps:
    - name: Checkout diffusers
      uses: actions/checkout@v3
      with:
        fetch-depth: 2
    - name: Install dependencies
      run: |
        apt-get update && apt-get install libsndfile1-dev libgl1 -y
        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
        python -m uv pip install -e [quality,test]
        python -m uv pip install 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:
        HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_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@v2
      with:
        name: flax_tpu_test_reports
        path: reports
    - name: Generate Report and Notify Channel
      if: always()
      run: |
        pip install slack_sdk tabulate
        python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY
  run_nightly_onnx_tests:
    name: Nightly ONNXRuntime CUDA tests on Ubuntu
    runs-on: docker-gpu
    container:
      image: diffusers/diffusers-onnxruntime-cuda
      options: --gpus 0 --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/
    steps:
    - name: Checkout diffusers
      uses: actions/checkout@v3
      with:
        fetch-depth: 2
    - name: NVIDIA-SMI
      run: nvidia-smi
    - name: Install dependencies
      run: |
        python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
        python -m uv pip install -e [quality,test]
        python -m uv pip install 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 ONNXRuntime CUDA tests
      env:
        HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }}
      run: |
        python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
          -s -v -k "Onnx" \
          --make-reports=tests_onnx_cuda \
          --report-log=tests_onnx_cuda.log \ 
          tests/
    - name: Failure short reports
      if: ${{ failure() }}
      run: |
        cat reports/tests_onnx_cuda_stats.txt
        cat reports/tests_onnx_cuda_failures_short.txt
    - name: Test suite reports artifacts
      if: ${{ always() }}
      uses: actions/upload-artifact@v2
      with:
        name: ${{ matrix.config.report }}_test_reports
        path: reports
    - name: Generate Report and Notify Channel
      if: always()
      run: |
        pip install slack_sdk tabulate
        python scripts/log_reports.py >> $GITHUB_STEP_SUMMARY
  run_nightly_tests_apple_m1:
    name: Nightly PyTorch MPS tests on MacOS
    runs-on: [ self-hosted, apple-m1 ]
@@ -35,10 +35,6 @@ jobs:
        run: |
          ruff check examples tests src utils scripts
          ruff format examples tests src utils scripts --check
      - name: Check if failure
        if: ${{ failure() }}
        run: |
          echo "Quality check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make style && make quality'" >> $GITHUB_STEP_SUMMARY
  check_repository_consistency:
    needs: check_code_quality
@@ -58,10 +54,6 @@ jobs:
          python utils/check_copies.py
          python utils/check_dummies.py
          make deps_table_check_updated
      - name: Check if failure
        if: ${{ failure() }}
        run: |
          echo "Repo consistency check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make fix-copies'" >> $GITHUB_STEP_SUMMARY
  run_fast_tests:
    needs: [check_code_quality, check_repository_consistency]
@@ -113,4 +105,4 @@ jobs:
        python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
          -s -v \
          --make-reports=tests_${{ matrix.config.report }} \
-          tests/lora/
+          tests/lora/test_lora_layers_peft.py
@@ -43,10 +43,6 @@ jobs:
        run: |
          ruff check examples tests src utils scripts
          ruff format examples tests src utils scripts --check
      - name: Check if failure
        if: ${{ failure() }}
        run: |
          echo "Quality check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make style && make quality'" >> $GITHUB_STEP_SUMMARY
  check_repository_consistency:
    needs: check_code_quality
@@ -66,10 +62,6 @@ jobs:
          python utils/check_copies.py
          python utils/check_dummies.py
          make deps_table_check_updated
      - name: Check if failure
        if: ${{ failure() }}
        run: |
          echo "Repo consistency check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make fix-copies'" >> $GITHUB_STEP_SUMMARY
  run_fast_tests:
    needs: [check_code_quality, check_repository_consistency]
@@ -21,7 +21,10 @@ env:
 jobs:
  setup_torch_cuda_pipeline_matrix:
    name: Setup Torch Pipelines CUDA Slow Tests Matrix
-    runs-on: ubuntu-latest
+    runs-on: [single-gpu, nvidia-gpu, t4, ci]
    container:
      image: diffusers/diffusers-pytorch-cpu # this is a CPU image, but we need it to fetch the matrix
      options: --shm-size "16gb" --ipc host
    outputs:
      pipeline_test_matrix: ${{ steps.fetch_pipeline_matrix.outputs.pipeline_test_matrix }}
    steps:
@@ -29,20 +32,24 @@ jobs:
        uses: actions/checkout@v3
        with:
          fetch-depth: 2
      - name: Set up Python
        uses: actions/setup-python@v4
        with:
          python-version: "3.8"
      - name: Install dependencies
        run: |
-          pip install -e .
+          apt-get update && apt-get install libsndfile1-dev libgl1 -y
-          pip install huggingface_hub
+          python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
          python -m uv pip install -e [quality,test]
          python -m uv pip install accelerate@git+https://github.com/huggingface/accelerate.git
      - name: Environment
        run: |
          python utils/print_env.py
      - name: Fetch Pipeline Matrix
        id: fetch_pipeline_matrix
        run: |
          matrix=$(python utils/fetch_torch_cuda_pipeline_test_matrix.py)
          echo $matrix
          echo "pipeline_test_matrix=$matrix" >> $GITHUB_OUTPUT
      - name: Pipeline Tests Artifacts
        if: ${{ always() }}
        uses: actions/upload-artifact@v2
@@ -52,9 +52,7 @@ jobs:
      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
-          pip install -U setuptools wheel twine
+          pip install -U setuptools wheel twine torch
          pip install -U torch --index-url https://download.pytorch.org/whl/cpu
          pip install -U transformers
      - name: Build the dist files
        run: python setup.py bdist_wheel && python setup.py sdist
@@ -19,16 +19,6 @@ authors:
    family-names: Rasul
  - given-names: Mishig
    family-names: Davaadorj
  - given-names: Dhruv
    family-names: Nair
  - given-names: Sayak
    family-names: Paul
  - given-names: Steven
    family-names: Liu
  - given-names: William
    family-names: Berman
  - given-names: Yiyi
    family-names: Xu
  - given-names: Thomas
    family-names: Wolf
 repository-code: 'https://github.com/huggingface/diffusers'
@@ -238,7 +238,7 @@ We also want to thank @heejkoo for the very helpful overview of papers, code and
 ```bibtex
@misc{von-platen-etal-2022-diffusers,
-  author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Dhruv Nair and Sayak Paul and William Berman and Yiyi Xu and Steven Liu and Thomas Wolf},
+  author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Thomas Wolf},
  title = {Diffusers: State-of-the-art diffusion models},
  year = {2022},
  publisher = {GitHub},
@@ -400,14 +400,14 @@
      title: DPMSolverSDEScheduler
    - local: api/schedulers/singlestep_dpm_solver
      title: DPMSolverSinglestepScheduler
    - local: api/schedulers/edm_multistep_dpm_solver
      title: EDMDPMSolverMultistepScheduler
    - local: api/schedulers/edm_euler
      title: EDMEulerScheduler
    - local: api/schedulers/euler_ancestral
      title: EulerAncestralDiscreteScheduler
    - local: api/schedulers/euler
      title: EulerDiscreteScheduler
    - local: api/schedulers/edm_euler
      title: EDMEulerScheduler
    - local: api/schedulers/edm_multistep_dpm_solver
      title: EDMDPMSolverMultistepScheduler
    - local: api/schedulers/heun
      title: HeunDiscreteScheduler
    - local: api/schedulers/ipndm
@@ -408,29 +408,6 @@ Make sure to check out the Schedulers [guide](../../using-diffusers/schedulers)
 </Tip>
 <table>
    <tr>
      <th align=center>Without FreeInit enabled</th>
      <th align=center>With FreeInit enabled</th>
    </tr>
    <tr>
        <td align=center>
          panda playing a guitar
          <br />
          <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-no-freeinit.gif"
              alt="panda playing a guitar"
              style="width: 300px;" />
        </td>
        <td align=center>
          panda playing a guitar
          <br/>
          <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/animatediff-freeinit.gif"
              alt="panda playing a guitar"
              style="width: 300px;" />
        </td>
    </tr>
 </table>
 ## Using AnimateLCM
 [AnimateLCM](https://animatelcm.github.io/) is a motion module checkpoint and an [LCM LoRA](https://huggingface.co/docs/diffusers/using-diffusers/inference_with_lcm_lora) that have been created using a consistency learning strategy that decouples the distillation of the image generation priors and the motion generation priors.
@@ -172,41 +172,3 @@ inpaint = StableDiffusionInpaintPipeline(**text2img.components)
 # now you can use text2img(...), img2img(...), inpaint(...) just like the call methods of each respective pipeline
 ```
 ### Create web demos using `gradio`
 The Stable Diffusion pipelines are automatically supported in [Gradio](https://github.com/gradio-app/gradio/), a library that makes creating beautiful and user-friendly machine learning apps on the web a breeze. First, make sure you have Gradio installed:
 ```
 pip install -U gradio
 ```
 Then, create a web demo around any Stable Diffusion-based pipeline. For example, you can create an image generation pipeline in a single line of code with Gradio's [`Interface.from_pipeline`](https://www.gradio.app/docs/interface#interface-from-pipeline) function:
 ```py
 from diffusers import StableDiffusionPipeline
 import gradio as gr
 pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
 gr.Interface.from_pipeline(pipe).launch()
 ```
 which opens an intuitive drag-and-drop interface in your browser:
 ![](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/gradio-panda.png)
 Similarly, you could create a demo for an image-to-image pipeline with:
 ```py
 from diffusers import StableDiffusionImg2ImgPipeline
 import gradio as gr
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
 gr.Interface.from_pipeline(pipe).launch()
 ```
 By default, the web demo runs on a local server. If you'd like to share it with others, you can generate a temporary public
 link by setting `share=True` in `launch()`. Or, you can host your demo on [Hugging Face Spaces](https://huggingface.co/spaces)https://huggingface.co/spaces for a permanent link. 
@@ -88,7 +88,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -54,7 +54,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -84,7 +84,7 @@ Many of the basic parameters are described in the [DreamBooth](dreambooth#script
 - `--freeze_model`: freezes the key and value parameters in the cross-attention layer; the default is `crossattn_kv`, but you can set it to `crossattn` to train all the parameters in the cross-attention layer
 - `--concepts_list`: to learn multiple concepts, provide a path to a JSON file containing the concepts
 - `--modifier_token`: a special word used to represent the learned concept
- `--initializer_token`: a special word used to initialize the embeddings of the `modifier_token`
+- `--initializer_token`:
 ### Prior preservation loss
@@ -67,7 +67,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -51,7 +51,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -89,7 +89,7 @@ The dataset preprocessing code and training loop are found in the [`main()`](htt
 As with the script parameters, a walkthrough of the training script is provided in the [Text-to-image](text2image#training-script) training guide. Instead, this guide takes a look at the InstructPix2Pix relevant parts of the script.
-The script begins by modifying the [number of input channels](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L445) in the first convolutional layer of the UNet to account for InstructPix2Pix's additional conditioning image:
+The script begins by modifing the [number of input channels](https://github.com/huggingface/diffusers/blob/64603389da01082055a901f2883c4810d1144edb/examples/instruct_pix2pix/train_instruct_pix2pix.py#L445) in the first convolutional layer of the UNet to account for InstructPix2Pix's additional conditioning image:
 ```py
 in_channels = 8
@@ -59,7 +59,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -53,7 +53,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -59,7 +59,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -53,7 +53,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -69,7 +69,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -67,7 +67,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -51,7 +51,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -53,7 +53,7 @@ accelerate config default
 Or if your environment doesn't support an interactive shell, like a notebook, you can use:
-```py
+```bash
 from accelerate.utils import write_basic_config
 write_basic_config()
@@ -45,7 +45,7 @@ Make sure to include the token `toy_face` in the prompt and then you can perform
 ```python
 prompt = "toy_face of a hacker with a hoodie"
-lora_scale = 0.9
+lora_scale= 0.9
 image = pipe(
    prompt, num_inference_steps=30, cross_attention_kwargs={"scale": lora_scale}, generator=torch.manual_seed(0)
 ).images[0]
@@ -114,7 +114,7 @@ To return to only using one adapter, use the [`~diffusers.loaders.UNet2DConditio
 pipe.set_adapters("toy")
 prompt = "toy_face of a hacker with a hoodie"
-lora_scale = 0.9
+lora_scale= 0.9
 image = pipe(
    prompt, num_inference_steps=30, cross_attention_kwargs={"scale": lora_scale}, generator=torch.manual_seed(0)
 ).images[0]
@@ -127,68 +127,11 @@ Or to disable all adapters entirely, use the [`~diffusers.loaders.UNet2DConditio
 pipe.disable_lora()
 prompt = "toy_face of a hacker with a hoodie"
 lora_scale= 0.9
 image = pipe(prompt, num_inference_steps=30, generator=torch.manual_seed(0)).images[0]
 image
 ```
 ![no-lora](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peft_integration/diffusers_peft_lora_inference_20_1.png)
 ### Customize adapters strength
 For even more customization, you can control how strongly the adapter affects each part of the pipeline. For this, pass a dictionary with the control strengths (called "scales") to [`~diffusers.loaders.UNet2DConditionLoadersMixin.set_adapters`].
 For example, here's how you can turn on the adapter for the `down` parts, but turn it off for the `mid` and `up` parts:
 ```python
 pipe.enable_lora()  # enable lora again, after we disabled it above
 prompt = "toy_face of a hacker with a hoodie, pixel art"
 adapter_weight_scales = { "unet": { "down": 1, "mid": 0, "up": 0} }
 pipe.set_adapters("pixel", adapter_weight_scales)
 image = pipe(prompt, num_inference_steps=30, generator=torch.manual_seed(0)).images[0]
 image
 ```
 ![block-lora-text-and-down](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peft_integration/diffusers_peft_lora_inference_block_down.png)
 Let's see how turning off the `down` part and turning on the `mid` and `up` part respectively changes the image.
 ```python
 adapter_weight_scales = { "unet": { "down": 0, "mid": 1, "up": 0} }
 pipe.set_adapters("pixel", adapter_weight_scales)
 image = pipe(prompt, num_inference_steps=30, generator=torch.manual_seed(0)).images[0]
 image
 ```
 ![block-lora-text-and-mid](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peft_integration/diffusers_peft_lora_inference_block_mid.png)
 ```python
 adapter_weight_scales = { "unet": { "down": 0, "mid": 0, "up": 1} }
 pipe.set_adapters("pixel", adapter_weight_scales)
 image = pipe(prompt, num_inference_steps=30, generator=torch.manual_seed(0)).images[0]
 image
 ```
 ![block-lora-text-and-up](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peft_integration/diffusers_peft_lora_inference_block_up.png)
 Looks cool!
 This is a really powerful feature. You can use it to control the adapter strengths down to per-transformer level. And you can even use it for multiple adapters.
 ```python
 adapter_weight_scales_toy = 0.5
 adapter_weight_scales_pixel = {
    "unet": {
        "down": 0.9,  # all transformers in the down-part will use scale 0.9
        # "mid"  # because, in this example, "mid" is not given, all transformers in the mid part will use the default scale 1.0
        "up": {
            "block_0": 0.6,  # all 3 transformers in the 0th block in the up-part will use scale 0.6
            "block_1": [0.4, 0.8, 1.0],  # the 3 transformers in the 1st block in the up-part will use scales 0.4, 0.8 and 1.0 respectively
        }
    }
 }
 pipe.set_adapters(["toy", "pixel"], [adapter_weight_scales_toy, adapter_weight_scales_pixel])
 image = pipe(prompt, num_inference_steps=30, generator=torch.manual_seed(0)).images[0]
 image
 ```
 ![block-lora-mixed](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/peft_integration/diffusers_peft_lora_inference_block_mixed.png)
 ## Manage active adapters
 You have attached multiple adapters in this tutorial, and if you're feeling a bit lost on what adapters have been attached to the pipeline's components, use the [`~diffusers.loaders.LoraLoaderMixin.get_active_adapters`] method to check the list of active adapters:
@@ -239,7 +239,5 @@ pipeline.to("cuda")
 prompt = "柴犬、カラフルアート"
 image = pipeline(prompt=prompt).images[0]
 ```
-> [!TIP]
+```
 > When using `trust_remote_code=True`, it is also strongly encouraged to pass a commit hash as a `revision` to make sure the author of the models did not update the code with some malicious new lines (unless you fully trust the authors of the models).
@@ -60,23 +60,6 @@ repo_id = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(repo_id)
 ```
 You can use the Space below to gauge the memory requirements of a pipeline you want to load beforehand without downloading the pipeline checkpoints:
 <div class="block dark:hidden">
 	<iframe 
        src="https://diffusers-compute-pipeline-size.hf.space?__theme=light"
        width="850"
        height="1600"
    ></iframe>
 </div>
 <div class="hidden dark:block">
    <iframe 
        src="https://diffusers-compute-pipeline-size.hf.space?__theme=dark"
        width="850"
        height="1600"
    ></iframe>
 </div>
 ### Local pipeline
 To load a diffusion pipeline locally, use [`git-lfs`](https://git-lfs.github.com/) to manually download the checkpoint (in this case, [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/runwayml/stable-diffusion-v1-5)) to your local disk. This creates a local folder, `./stable-diffusion-v1-5`, on your disk:
@@ -153,43 +153,18 @@ image
    <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/load_attn_proc.png" />
 </div>
 <Tip>
 For both [`~loaders.LoraLoaderMixin.load_lora_weights`] and [`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`], you can pass the `cross_attention_kwargs={"scale": 0.5}` parameter to adjust how much of the LoRA weights to use. A value of `0` is the same as only using the base model weights, and a value of `1` is equivalent to using the fully finetuned LoRA.
 </Tip>
 To unload the LoRA weights, use the [`~loaders.LoraLoaderMixin.unload_lora_weights`] method to discard the LoRA weights and restore the model to its original weights:
 ```py
 pipeline.unload_lora_weights()
 ```
 ### Adjust LoRA weight scale
 For both [`~loaders.LoraLoaderMixin.load_lora_weights`] and [`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`], you can pass the `cross_attention_kwargs={"scale": 0.5}` parameter to adjust how much of the LoRA weights to use. A value of `0` is the same as only using the base model weights, and a value of `1` is equivalent to using the fully finetuned LoRA.
 For more granular control on the amount of LoRA weights used per layer, you can use [`~loaders.LoraLoaderMixin.set_adapters`] and pass a dictionary specifying by how much to scale the weights in each layer by.
 ```python
 pipe = ... # create pipeline
 pipe.load_lora_weights(..., adapter_name="my_adapter") 
 scales = {
    "text_encoder": 0.5,
    "text_encoder_2": 0.5,  # only usable if pipe has a 2nd text encoder
    "unet": {
        "down": 0.9,  # all transformers in the down-part will use scale 0.9
        # "mid"  # in this example "mid" is not given, therefore all transformers in the mid part will use the default scale 1.0
        "up": {
            "block_0": 0.6,  # all 3 transformers in the 0th block in the up-part will use scale 0.6
            "block_1": [0.4, 0.8, 1.0],  # the 3 transformers in the 1st block in the up-part will use scales 0.4, 0.8 and 1.0 respectively
        }
    }
 }
 pipe.set_adapters("my_adapter", scales)
 ```
 This also works with multiple adapters - see [this guide](https://huggingface.co/docs/diffusers/tutorials/using_peft_for_inference#customize-adapters-strength) for how to do it.
 <Tip warning={true}>
 Currently, [`~loaders.LoraLoaderMixin.set_adapters`] only supports scaling attention weights. If a LoRA has other parts (e.g., resnets or down-/upsamplers), they will keep a scale of 1.0.
 </Tip>
 ### Kohya and TheLastBen
 Other popular LoRA trainers from the community include those by [Kohya](https://github.com/kohya-ss/sd-scripts/) and [TheLastBen](https://github.com/TheLastBen/fast-stable-diffusion). These trainers create different LoRA checkpoints than those trained by 🤗 Diffusers, but they can still be loaded in the same way.
@@ -86,7 +86,7 @@ Video generation is very memory intensive because you're essentially generating
 + frames = pipe(image, decode_chunk_size=2, generator=generator, num_frames=25).frames[0]
 ```
-Using all these tricks together should lower the memory requirement to less than 8GB VRAM.
+Using all these tricks togethere should lower the memory requirement to less than 8GB VRAM.
 ## Micro-conditioning
@@ -355,7 +355,7 @@ unet_traced = torch.jit.load("unet_traced.pt")
 class TracedUNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
-        self.in_channels = pipe.unet.config.in_channels
+        self.in_channels = pipe.unet.in_channels
        self.device = pipe.unet.device
    def forward(self, latent_model_input, t, encoder_hidden_states):
@@ -210,7 +210,7 @@ Stable Diffusion 은 text-to-image *latent diffusion* 모델입니다. latent di
 ```py
 >>> latents = torch.randn(
-...     (batch_size, unet.config.in_channels, height // 8, width // 8),
+...     (batch_size, unet.in_channels, height // 8, width // 8),
 ...     generator=generator,
 ...     device=torch_device,
 ... )
@@ -42,7 +42,7 @@ Training examples show how to pretrain or fine-tune diffusion models for a varie
 | [**Dreambooth**](./dreambooth) | ✅ | - | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_dreambooth_training.ipynb)
 | [**ControlNet**](./controlnet) | ✅ | ✅ | -
 | [**InstructPix2Pix**](./instruct_pix2pix) | ✅ | ✅ | -
-| [**Reinforcement Learning for Control**](./reinforcement_learning)                    | - | - | coming soon.
+| [**Reinforcement Learning for Control**](https://github.com/huggingface/diffusers/blob/main/examples/reinforcement_learning/run_diffusers_locomotion.py)                    | - | - | coming soon.
 ## Community
@@ -308,6 +308,6 @@ accelerate launch train_dreambooth_lora_sdxl_advanced.py \
 Check out [these recommended practices](https://huggingface.co/blog/sdxl_lora_advanced_script#additional-good-practices)
 ## Running on Colab Notebook
-Check out [this notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_Dreambooth_LoRA_advanced_example.ipynb).
+Check out [this notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_DreamBooth_LoRA_advanced_example.ipynb). 
 to train using the advanced features (including pivotal tuning), and [this notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_DreamBooth_LoRA_.ipynb) to train on a free colab, using some of the advanced features (excluding pivotal tuning)
@@ -70,7 +70,7 @@ from diffusers.utils.import_utils import is_xformers_available
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -656,6 +656,7 @@ def parse_args(input_args=None):
    )
    parser.add_argument(
        "--use_dora",
        type=bool,
        action="store_true",
        default=False,
        help=(
@@ -78,7 +78,7 @@ from diffusers.utils.torch_utils import is_compiled_module
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -1,16 +1,13 @@
-# Community Pipeline Examples
+# Community Examples
 > **For more information about community pipelines, please have a look at [this issue](https://github.com/huggingface/diffusers/issues/841).**
-**Community pipeline** examples consist pipelines that have been added by the community.
+**Community** examples consist of both inference and training examples that have been added by the community.
-Please have a look at the following tables to get an overview of all community examples. Click on the **Code Example** to get a copy-and-paste ready code example that you can try out.
+Please have a look at the following table to get an overview of all community examples. Click on the **Code Example** to get a copy-and-paste ready code example that you can try out.
-If a community pipeline doesn't work as expected, please open an issue and ping the author on it.
+If a community doesn't work as expected, please open an issue and ping the author on it.
 Please also check out our [Community Scripts](https://github.com/huggingface/diffusers/blob/main/examples/community/README_community_scripts.md) examples for tips and tricks that you can use with diffusers without having to run a community pipeline.
 | Example                                                                                                                               | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              | Code Example                                                                              | Colab                                                                                                                                                                                                              |                                                        Author |
 |:--------------------------------------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------:|
 | HD-Painter                                                                                                                            | [HD-Painter](https://github.com/Picsart-AI-Research/HD-Painter) enables prompt-faithfull and high resolution (up to 2k) image inpainting upon any diffusion-based image inpainting method.                                                                                                                                                                                                                                                                                                               | [HD-Painter](#hd-painter)                                                                 | [![Hugging Face Space](https://img.shields.io/badge/🤗%20Hugging%20Face-Space-yellow)](https://huggingface.co/spaces/PAIR/HD-Painter)                                                                              | [Manukyan Hayk](https://github.com/haikmanukyan) and [Sargsyan Andranik](https://github.com/AndranikSargsyan) |
 | Marigold Monocular Depth Estimation                                                                                                   | A universal monocular depth estimator, utilizing Stable Diffusion, delivering sharp predictions in the wild. (See the [project page](https://marigoldmonodepth.github.io) and [full codebase](https://github.com/prs-eth/marigold) for more details.)                                                                                                                                                                                                                                                        | [Marigold Depth Estimation](#marigold-depth-estimation)                                   | [![Hugging Face Space](https://img.shields.io/badge/🤗%20Hugging%20Face-Space-yellow)](https://huggingface.co/spaces/toshas/marigold) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/12G8reD13DdpMie5ZQlaFNo2WCGeNUH-u?usp=sharing) | [Bingxin Ke](https://github.com/markkua) and [Anton Obukhov](https://github.com/toshas) |
 | LLM-grounded Diffusion (LMD+)                                                                                                         | LMD greatly improves the prompt following ability of text-to-image generation models by introducing an LLM as a front-end prompt parser and layout planner. [Project page.](https://llm-grounded-diffusion.github.io/) [See our full codebase (also with diffusers).](https://github.com/TonyLianLong/LLM-groundedDiffusion)                                                                                                                                                                                                                                                                                                                                                                                                                                   | [LLM-grounded Diffusion (LMD+)](#llm-grounded-diffusion)                             | [Huggingface Demo](https://huggingface.co/spaces/longlian/llm-grounded-diffusion) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1SXzMSeAB-LJYISb2yrUOdypLz4OYWUKj) |                [Long (Tony) Lian](https://tonylian.com/) |
 | CLIP Guided Stable Diffusion                                                                                                          | Doing CLIP guidance for text to image generation with Stable Diffusion                                                                                                                                                                                                                                                                                                                                                                                                                                   | [CLIP Guided Stable Diffusion](#clip-guided-stable-diffusion)                             | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/CLIP_Guided_Stable_diffusion_with_diffusers.ipynb) |                [Suraj Patil](https://github.com/patil-suraj/) |
@@ -76,48 +73,6 @@ pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", custo
 ## Example usages
 ### HD-Painter
 Implementation of [HD-Painter: High-Resolution and Prompt-Faithful Text-Guided Image Inpainting with Diffusion Models](https://arxiv.org/abs/2312.14091).
 ![teaser-img](https://raw.githubusercontent.com/Picsart-AI-Research/HD-Painter/main/__assets__/github/teaser.jpg)
 The abstract from the paper is:
 Recent progress in text-guided image inpainting, based on the unprecedented success of text-to-image diffusion models, has led to exceptionally realistic and visually plausible results.
 However, there is still significant potential for improvement in current text-to-image inpainting models, particularly in better aligning the inpainted area with user prompts and performing high-resolution inpainting.
 Therefore, in this paper we introduce _HD-Painter_, a completely **training-free** approach that **accurately follows to prompts** and coherently **scales to high-resolution** image inpainting.
 To this end, we design the _Prompt-Aware Introverted Attention (PAIntA)_ layer enhancing self-attention scores by prompt information and resulting in better text alignment generations.
 To further improve the prompt coherence we introduce the _Reweighting Attention Score Guidance (RASG)_ mechanism seamlessly integrating a post-hoc sampling strategy into general form of DDIM to prevent out-of-distribution latent shifts.
 Moreover, HD-Painter allows extension to larger scales by introducing a specialized super-resolution technique customized for inpainting, enabling the completion of missing regions in images of up to 2K resolution. 
 Our experiments demonstrate that HD-Painter surpasses existing state-of-the-art approaches qualitatively and quantitatively, achieving an impressive generation accuracy improvement of **61.4** vs **51.9**. 
 We will make the codes publicly available.
 You can find additional information about Text2Video-Zero in the [paper](https://arxiv.org/abs/2312.14091) or the [original codebase](https://github.com/Picsart-AI-Research/HD-Painter).
 #### Usage example
 ```python
 import torch
 from diffusers import DiffusionPipeline, DDIMScheduler
 from diffusers.utils import load_image, make_image_grid
 pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting",
    custom_pipeline="hd_painter"
 )
 pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 prompt = "wooden boat"
 init_image = load_image("https://raw.githubusercontent.com/Picsart-AI-Research/HD-Painter/main/__assets__/samples/images/2.jpg")
 mask_image = load_image("https://raw.githubusercontent.com/Picsart-AI-Research/HD-Painter/main/__assets__/samples/masks/2.png")
 image = pipe (prompt, init_image, mask_image, use_rasg = True, use_painta = True, generator=torch.manual_seed(12345)).images[0]
 make_image_grid([init_image, mask_image, image], rows=1, cols=3)
 ```
 ### Marigold Depth Estimation
 Marigold is a universal monocular depth estimator that delivers accurate and sharp predictions in the wild. Based on Stable Diffusion, it is trained exclusively with synthetic depth data and excels in zero-shot adaptation to real-world imagery. This pipeline is an official implementation of the inference process. More details can be found on our [project page](https://marigoldmonodepth.github.io) and [full codebase](https://github.com/prs-eth/marigold) (also implemented with diffusers).
@@ -128,25 +83,14 @@ This depth estimation pipeline processes a single input image through multiple d
 ```python
 import numpy as np
 import torch
 from PIL import Image
 from diffusers import DiffusionPipeline
 from diffusers.utils import load_image
 # Original DDIM version (higher quality)
 pipe = DiffusionPipeline.from_pretrained(
-    "prs-eth/marigold-v1-0",
+    "Bingxin/Marigold",
    custom_pipeline="marigold_depth_estimation"
    # torch_dtype=torch.float16,                # (optional) Run with half-precision (16-bit float).
    # variant="fp16",                           # (optional) Use with `torch_dtype=torch.float16`, to directly load fp16 checkpoint
 )
 # (New) LCM version (faster speed)
 pipe = DiffusionPipeline.from_pretrained(
    "prs-eth/marigold-lcm-v1-0",
    custom_pipeline="marigold_depth_estimation"
    # torch_dtype=torch.float16,                # (optional) Run with half-precision (16-bit float).
    # variant="fp16",                           # (optional) Use with `torch_dtype=torch.float16`, to directly load fp16 checkpoint
 )
 pipe.to("cuda")
@@ -155,21 +99,12 @@ img_path_or_url = "https://share.phys.ethz.ch/~pf/bingkedata/marigold/pipeline_e
 image: Image.Image = load_image(img_path_or_url)
 pipeline_output = pipe(
-    image,                    # Input image.
+    image,                  # Input image.
    # ----- recommended setting for DDIM version -----
    # denoising_steps=10,     # (optional) Number of denoising steps of each inference pass. Default: 10.
    # ensemble_size=10,       # (optional) Number of inference passes in the ensemble. Default: 10.
    # ------------------------------------------------
    # ----- recommended setting for LCM version ------
    # denoising_steps=4,
    # ensemble_size=5,
    # -------------------------------------------------
    # processing_res=768,     # (optional) Maximum resolution of processing. If set to 0: will not resize at all. Defaults to 768.
    # match_input_res=True,   # (optional) Resize depth prediction to match input resolution.
    # batch_size=0,           # (optional) Inference batch size, no bigger than `num_ensemble`. If set to 0, the script will automatically decide the proper batch size. Defaults to 0.
    # seed=2024,              # (optional) Random seed can be set to ensure additional reproducibility. Default: None (unseeded). Note: forcing --batch_size 1 helps to increase reproducibility. To ensure full reproducibility, deterministic mode needs to be used.
    # color_map="Spectral",   # (optional) Colormap used to colorize the depth map. Defaults to "Spectral". Set to `None` to skip colormap generation.
    # show_progress_bar=True, # (optional) If true, will show progress bars of the inference progress.
 )
@@ -998,7 +933,7 @@ image = pipe(prompt, generator=generator, num_inference_steps=50).images[0]
 ### Checkpoint Merger Pipeline
 Based on the AUTOMATIC1111/webui for checkpoint merging. This is a custom pipeline that merges upto 3 pretrained model checkpoints as long as they are in the HuggingFace model_index.json format.
-The checkpoint merging is currently memory intensive as it modifies the weights of a DiffusionPipeline object in place. Expect at least 13GB RAM Usage on Kaggle GPU kernels and
+The checkpoint merging is currently memory intensive as it modifies the weights of a DiffusionPipeline object in place. Expect atleast 13GB RAM Usage on Kaggle GPU kernels and
 on colab you might run out of the 12GB memory even while merging two checkpoints.
 Usage:-
@@ -1952,7 +1887,7 @@ In the above code, the `prompt2` is appended to the `prompt`, which is more than
 For more results, checkout [PR #6114](https://github.com/huggingface/diffusers/pull/6114).
-### Example Images Mixing (with CoCa)
+## Example Images Mixing (with CoCa)
 ```python
 import requests
 from io import BytesIO
@@ -2999,7 +2934,7 @@ pipe(prompt =prompt, rp_args = rp_args)
 The Pipeline supports `compel` syntax. Input prompts using the `compel` structure will be automatically applied and processed.
-### Diffusion Posterior Sampling Pipeline
+## Diffusion Posterior Sampling Pipeline
 * Reference paper
    ```
    @article{chung2022diffusion,
@@ -3806,80 +3741,3 @@ onestep_image = pipe(prompt, num_inference_steps=1).images[0]
 # Multistep sampling
 multistep_image = pipe(prompt, num_inference_steps=4).images[0]
 ```
 # Perturbed-Attention Guidance
 [Project](https://ku-cvlab.github.io/Perturbed-Attention-Guidance/) / [arXiv](https://arxiv.org/abs/2403.17377) / [GitHub](https://github.com/KU-CVLAB/Perturbed-Attention-Guidance)
 This implementation is based on [Diffusers](https://huggingface.co/docs/diffusers/index). StableDiffusionPAGPipeline is a modification of StableDiffusionPipeline to support Perturbed-Attention Guidance (PAG).
 ## Example Usage
 ```
 import os
 import torch
 from accelerate.utils import set_seed
 from diffusers import StableDiffusionPipeline
 from diffusers.utils import load_image, make_image_grid
 from diffusers.utils.torch_utils import randn_tensor
 pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    custom_pipeline="hyoungwoncho/sd_perturbed_attention_guidance",
    torch_dtype=torch.float16
 )
 device="cuda"
 pipe = pipe.to(device)
 pag_scale = 5.0
 pag_applied_layers_index = ['m0']
 batch_size = 4
 seed=10
 base_dir = "./results/"
 grid_dir = base_dir + "/pag" + str(pag_scale) + "/"
 if not os.path.exists(grid_dir):
    os.makedirs(grid_dir)
 set_seed(seed)
 latent_input = randn_tensor(shape=(batch_size,4,64,64),generator=None, device=device, dtype=torch.float16)
 output_baseline = pipe(
    "",
    width=512,
    height=512,
    num_inference_steps=50,
    guidance_scale=0.0,
    pag_scale=0.0,
    pag_applied_layers_index=pag_applied_layers_index,
    num_images_per_prompt=batch_size,
    latents=latent_input
 ).images
 output_pag = pipe(
    "",
    width=512,
    height=512,
    num_inference_steps=50,
    guidance_scale=0.0,
    pag_scale=5.0,
    pag_applied_layers_index=pag_applied_layers_index,
    num_images_per_prompt=batch_size,
    latents=latent_input
 ).images
 grid_image = make_image_grid(output_baseline + output_pag, rows=2, cols=batch_size)
 grid_image.save(grid_dir + "sample.png")
 ```
 ## PAG Parameters
 pag_scale : gudiance scale of PAG (ex: 5.0)
 pag_applied_layers_index : index of the layer to apply perturbation (ex: ['m0'])
@@ -1,232 +0,0 @@
 # Community Scripts
 **Community scripts** consist of inference examples using Diffusers pipelines that have been added by the community. 
 Please have a look at the following table to get an overview of all community examples. Click on the **Code Example** to get a copy-and-paste code example that you can try out.
 If a community script doesn't work as expected, please open an issue and ping the author on it.
 | Example                                                                                                                               | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              | Code Example                                                                              | Colab                                                                                                                                                                                                              |                                                        Author |
 |:--------------------------------------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------------------------------------------------------------------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------:|
 | Using IP-Adapter with negative noise                                                                                                  | Using negative noise with IP-adapter to better control the generation (see the [original post](https://github.com/huggingface/diffusers/discussions/7167) on the forum for more details)                                                                                                                                                                                                                                                    | [IP-Adapter Negative Noise](#ip-adapter-negative-noise)                                   | | [Álvaro Somoza](https://github.com/asomoza)|
 | asymmetric tiling                                                                                                  |configure seamless image tiling independently for the X and Y axes                                                                                                                                                                                                      | [Asymmetric Tiling](#asymmetric-tiling )                                   | | [alexisrolland](https://github.com/alexisrolland)|
 ## Example usages
 ### IP Adapter Negative Noise
 Diffusers pipelines are fully integrated with IP-Adapter, which allows you to prompt the diffusion model with an image. However, it does not support negative image prompts (there is no `negative_ip_adapter_image` argument) the same way it supports negative text prompts. When you pass an `ip_adapter_image,` it will create a zero-filled tensor as a negative image. This script shows you how to create a negative noise from `ip_adapter_image` and use it to significantly improve the generation quality while preserving the composition of images.
 [cubiq](https://github.com/cubiq) initially developed this feature in his [repository](https://github.com/cubiq/ComfyUI_IPAdapter_plus). The community script was contributed by [asomoza](https://github.com/Somoza). You can find more details about this experimentation [this discussion](https://github.com/huggingface/diffusers/discussions/7167)
 IP-Adapter without negative noise
 |source|result|
 |---|---|
 |![20240229150812](https://github.com/huggingface/diffusers/assets/5442875/901d8bd8-7a59-4fe7-bda1-a0e0d6c7dffd)|![20240229163923_normal](https://github.com/huggingface/diffusers/assets/5442875/3432e25a-ece6-45f4-a3f4-fca354f40b5b)|
 IP-Adapter with negative noise
 |source|result|
 |---|---|
 |![20240229150812](https://github.com/huggingface/diffusers/assets/5442875/901d8bd8-7a59-4fe7-bda1-a0e0d6c7dffd)|![20240229163923](https://github.com/huggingface/diffusers/assets/5442875/736fd15a-36ba-40c0-a7d8-6ec1ac26f788)|
 ```python
 import torch
 from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, StableDiffusionXLPipeline
 from diffusers.models import ImageProjection
 from diffusers.utils import load_image
 def encode_image(
    image_encoder,
    feature_extractor,
    image,
    device,
    num_images_per_prompt,
    output_hidden_states=None,
    negative_image=None,
 ):
    dtype = next(image_encoder.parameters()).dtype
    if not isinstance(image, torch.Tensor):
        image = feature_extractor(image, return_tensors="pt").pixel_values
    image = image.to(device=device, dtype=dtype)
    if output_hidden_states:
        image_enc_hidden_states = image_encoder(image, output_hidden_states=True).hidden_states[-2]
        image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
        if negative_image is None:
            uncond_image_enc_hidden_states = image_encoder(
                torch.zeros_like(image), output_hidden_states=True
            ).hidden_states[-2]
        else:
            if not isinstance(negative_image, torch.Tensor):
                negative_image = feature_extractor(negative_image, return_tensors="pt").pixel_values
            negative_image = negative_image.to(device=device, dtype=dtype)
            uncond_image_enc_hidden_states = image_encoder(negative_image, output_hidden_states=True).hidden_states[-2]
        uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
        return image_enc_hidden_states, uncond_image_enc_hidden_states
    else:
        image_embeds = image_encoder(image).image_embeds
        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
        uncond_image_embeds = torch.zeros_like(image_embeds)
        return image_embeds, uncond_image_embeds
@torch.no_grad()
 def prepare_ip_adapter_image_embeds(
    unet,
    image_encoder,
    feature_extractor,
    ip_adapter_image,
    do_classifier_free_guidance,
    device,
    num_images_per_prompt,
    ip_adapter_negative_image=None,
 ):
    if not isinstance(ip_adapter_image, list):
        ip_adapter_image = [ip_adapter_image]
    if len(ip_adapter_image) != len(unet.encoder_hid_proj.image_projection_layers):
        raise ValueError(
            f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
        )
    image_embeds = []
    for single_ip_adapter_image, image_proj_layer in zip(
        ip_adapter_image, unet.encoder_hid_proj.image_projection_layers
    ):
        output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
        single_image_embeds, single_negative_image_embeds = encode_image(
            image_encoder,
            feature_extractor,
            single_ip_adapter_image,
            device,
            1,
            output_hidden_state,
            negative_image=ip_adapter_negative_image,
        )
        single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
        single_negative_image_embeds = torch.stack([single_negative_image_embeds] * num_images_per_prompt, dim=0)
        if do_classifier_free_guidance:
            single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
            single_image_embeds = single_image_embeds.to(device)
        image_embeds.append(single_image_embeds)
    return image_embeds
 vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=torch.float16,
 ).to("cuda")
 pipeline = StableDiffusionXLPipeline.from_pretrained(
    "RunDiffusion/Juggernaut-XL-v9",
    torch_dtype=torch.float16,
    vae=vae,
    variant="fp16",
 ).to("cuda")
 pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
 pipeline.scheduler.config.use_karras_sigmas = True
 pipeline.load_ip_adapter(
    "h94/IP-Adapter",
    subfolder="sdxl_models",
    weight_name="ip-adapter-plus_sdxl_vit-h.safetensors",
    image_encoder_folder="models/image_encoder",
 )
 pipeline.set_ip_adapter_scale(0.7)
 ip_image = load_image("source.png")
 negative_ip_image = load_image("noise.png")
 image_embeds = prepare_ip_adapter_image_embeds(
    unet=pipeline.unet,
    image_encoder=pipeline.image_encoder,
    feature_extractor=pipeline.feature_extractor,
    ip_adapter_image=[[ip_image]],
    do_classifier_free_guidance=True,
    device="cuda",
    num_images_per_prompt=1,
    ip_adapter_negative_image=negative_ip_image,
 )
 prompt = "cinematic photo of a cyborg in the city, 4k, high quality, intricate, highly detailed"
 negative_prompt = "blurry, smooth, plastic"
 image = pipeline(
    prompt=prompt,
    negative_prompt=negative_prompt,
    ip_adapter_image_embeds=image_embeds,
    guidance_scale=6.0,
    num_inference_steps=25,
    generator=torch.Generator(device="cpu").manual_seed(1556265306),
 ).images[0]
 image.save("result.png")
 ```
 ### Asymmetric Tiling
 Stable Diffusion is not trained to generate seamless textures. However, you can use this simple script to add tiling to your generation. This script is contributed by [alexisrolland](https://github.com/alexisrolland). See more details in the [this issue](https://github.com/huggingface/diffusers/issues/556)
 |Generated|Tiled|
 |---|---|
 |![20240313003235_573631814](https://github.com/huggingface/diffusers/assets/5442875/eca174fb-06a4-464e-a3a7-00dbb024543e)|![wall](https://github.com/huggingface/diffusers/assets/5442875/b4aa774b-2a6a-4316-a8eb-8f30b5f4d024)|
 ```py
 import torch
 from typing import Optional
 from diffusers import StableDiffusionPipeline
 from diffusers.models.lora import LoRACompatibleConv
 def seamless_tiling(pipeline, x_axis, y_axis):
    def asymmetric_conv2d_convforward(self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None):
        self.paddingX = (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0)
        self.paddingY = (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3])
        working = torch.nn.functional.pad(input, self.paddingX, mode=x_mode)
        working = torch.nn.functional.pad(working, self.paddingY, mode=y_mode)
        return torch.nn.functional.conv2d(working, weight, bias, self.stride, torch.nn.modules.utils._pair(0), self.dilation, self.groups)
    x_mode = 'circular' if x_axis else 'constant'
    y_mode = 'circular' if y_axis else 'constant'
    targets = [pipeline.vae, pipeline.text_encoder, pipeline.unet]
    convolution_layers = []
    for target in targets:
        for module in target.modules():
            if isinstance(module, torch.nn.Conv2d):
                convolution_layers.append(module)
    for layer in convolution_layers:
        if isinstance(layer, LoRACompatibleConv) and layer.lora_layer is None:
            layer.lora_layer = lambda * x: 0
        layer._conv_forward = asymmetric_conv2d_convforward.__get__(layer, torch.nn.Conv2d)
    return pipeline
 pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16, use_safetensors=True)
 pipeline.enable_model_cpu_offload()
 prompt = ["texture of a red brick wall"]
 seed = 123456
 generator = torch.Generator(device='cuda').manual_seed(seed)
 pipeline = seamless_tiling(pipeline=pipeline, x_axis=True, y_axis=True)
 image = pipeline(
    prompt=prompt,
    width=512,
    height=512,
    num_inference_steps=20,
    guidance_scale=7,
    num_images_per_prompt=1,
    generator=generator
 ).images[0]
 seamless_tiling(pipeline=pipeline, x_axis=False, y_axis=False)
 torch.cuda.empty_cache()
 image.save('image.png')
 ```
@@ -103,7 +103,7 @@ class CheckpointMergerPipeline(DiffusionPipeline):
        print(f"Combining with alpha={alpha}, interpolation mode={interp}")
        checkpoint_count = len(pretrained_model_name_or_path_list)
-        # Ignore result from model_index_json comparison of the two checkpoints
+        # Ignore result from model_index_json comparision of the two checkpoints
        force = kwargs.pop("force", False)
        # If less than 2 checkpoints, nothing to merge. If more than 3, not supported for now.
@@ -217,7 +217,7 @@ class CheckpointMergerPipeline(DiffusionPipeline):
                        ]
                        checkpoint_path_2 = files[0] if len(files) > 0 else None
                # For an attr if both checkpoint_path_1 and 2 are None, ignore.
-                # If at least one is present, deal with it according to interp method, of course only if the state_dict keys match.
+                # If atleast one is present, deal with it according to interp method, of course only if the state_dict keys match.
                if checkpoint_path_1 is None and checkpoint_path_2 is None:
                    print(f"Skipping {attr}: not present in 2nd or 3d model")
                    continue
@@ -1,994 +0,0 @@
 import math
 import numbers
 from typing import Any, Callable, Dict, List, Optional, Union
 import torch
 import torch.nn.functional as F
 from torch import nn
 from diffusers.image_processor import PipelineImageInput
 from diffusers.models import AsymmetricAutoencoderKL, ImageProjection
 from diffusers.models.attention_processor import Attention, AttnProcessor
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import (
    StableDiffusionInpaintPipeline,
    retrieve_timesteps,
 )
 from diffusers.utils import deprecate
 class RASGAttnProcessor:
    def __init__(self, mask, token_idx, scale_factor):
        self.attention_scores = None  # Stores the last output of the similarity matrix here. Each layer will get its own RASGAttnProcessor assigned
        self.mask = mask
        self.token_idx = token_idx
        self.scale_factor = scale_factor
        self.mask_resoltuion = mask.shape[-1] * mask.shape[-2]  # 64 x 64 if the image is 512x512
    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        temb: Optional[torch.FloatTensor] = None,
        scale: float = 1.0,
    ) -> torch.Tensor:
        # Same as the default AttnProcessor up untill the part where similarity matrix gets saved
        downscale_factor = self.mask_resoltuion // hidden_states.shape[1]
        residual = hidden_states
        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)
        input_ndim = hidden_states.ndim
        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
        )
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
        query = attn.to_q(hidden_states)
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
        query = attn.head_to_batch_dim(query)
        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)
        # Automatically recognize the resolution and save the attention similarity values
        # We need to use the values before the softmax function, hence the rewritten get_attention_scores function.
        if downscale_factor == self.scale_factor**2:
            self.attention_scores = get_attention_scores(attn, query, key, attention_mask)
            attention_probs = self.attention_scores.softmax(dim=-1)
            attention_probs = attention_probs.to(query.dtype)
        else:
            attention_probs = attn.get_attention_scores(query, key, attention_mask)  # Original code
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)
        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)
        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
        if attn.residual_connection:
            hidden_states = hidden_states + residual
        hidden_states = hidden_states / attn.rescale_output_factor
        return hidden_states
 class PAIntAAttnProcessor:
    def __init__(self, transformer_block, mask, token_idx, do_classifier_free_guidance, scale_factors):
        self.transformer_block = transformer_block  # Stores the parent transformer block.
        self.mask = mask
        self.scale_factors = scale_factors
        self.do_classifier_free_guidance = do_classifier_free_guidance
        self.token_idx = token_idx
        self.shape = mask.shape[2:]
        self.mask_resoltuion = mask.shape[-1] * mask.shape[-2]  # 64 x 64
        self.default_processor = AttnProcessor()
    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        temb: Optional[torch.FloatTensor] = None,
        scale: float = 1.0,
    ) -> torch.Tensor:
        # Automatically recognize the resolution of the current attention layer and resize the masks accordingly
        downscale_factor = self.mask_resoltuion // hidden_states.shape[1]
        mask = None
        for factor in self.scale_factors:
            if downscale_factor == factor**2:
                shape = (self.shape[0] // factor, self.shape[1] // factor)
                mask = F.interpolate(self.mask, shape, mode="bicubic")  # B, 1, H, W
                break
        if mask is None:
            return self.default_processor(attn, hidden_states, encoder_hidden_states, attention_mask, temb, scale)
        # STARTS HERE
        residual = hidden_states
        # Save the input hidden_states for later use
        input_hidden_states = hidden_states
        # ================================================== #
        # =============== SELF ATTENTION 1 ================= #
        # ================================================== #
        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)
        input_ndim = hidden_states.ndim
        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
        )
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
        query = attn.to_q(hidden_states)
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
        query = attn.head_to_batch_dim(query)
        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)
        # self_attention_probs = attn.get_attention_scores(query, key, attention_mask) # We can't use post-softmax attention scores in this case
        self_attention_scores = get_attention_scores(
            attn, query, key, attention_mask
        )  # The custom function returns pre-softmax probabilities
        self_attention_probs = self_attention_scores.softmax(
            dim=-1
        )  # Manually compute the probabilities here, the scores will be reused in the second part of PAIntA
        self_attention_probs = self_attention_probs.to(query.dtype)
        hidden_states = torch.bmm(self_attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)
        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)
        # x = x + self.attn1(self.norm1(x))
        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
        if attn.residual_connection:  # So many residuals everywhere
            hidden_states = hidden_states + residual
        self_attention_output_hidden_states = hidden_states / attn.rescale_output_factor
        # ================================================== #
        # ============ BasicTransformerBlock =============== #
        # ================================================== #
        # We use a hack by running the code from the BasicTransformerBlock that is between Self and Cross attentions here
        # The other option would've been modifying the BasicTransformerBlock and adding this functionality here.
        # I assumed that changing the BasicTransformerBlock would have been a bigger deal and decided to use this hack isntead.
        # The SelfAttention block recieves the normalized latents from the BasicTransformerBlock,
        # But the residual of the output is the non-normalized version.
        # Therefore we unnormalize the input hidden state here
        unnormalized_input_hidden_states = (
            input_hidden_states + self.transformer_block.norm1.bias
        ) * self.transformer_block.norm1.weight
        # TODO: return if neccessary
        # if self.use_ada_layer_norm_zero:
        #     attn_output = gate_msa.unsqueeze(1) * attn_output
        # elif self.use_ada_layer_norm_single:
        #     attn_output = gate_msa * attn_output
        transformer_hidden_states = self_attention_output_hidden_states + unnormalized_input_hidden_states
        if transformer_hidden_states.ndim == 4:
            transformer_hidden_states = transformer_hidden_states.squeeze(1)
        # TODO: return if neccessary
        # 2.5 GLIGEN Control
        # if gligen_kwargs is not None:
        #     transformer_hidden_states = self.fuser(transformer_hidden_states, gligen_kwargs["objs"])
        # NOTE: we experimented with using GLIGEN and HDPainter together, the results were not that great
        # 3. Cross-Attention
        if self.transformer_block.use_ada_layer_norm:
            # transformer_norm_hidden_states = self.transformer_block.norm2(transformer_hidden_states, timestep)
            raise NotImplementedError()
        elif self.transformer_block.use_ada_layer_norm_zero or self.transformer_block.use_layer_norm:
            transformer_norm_hidden_states = self.transformer_block.norm2(transformer_hidden_states)
        elif self.transformer_block.use_ada_layer_norm_single:
            # For PixArt norm2 isn't applied here:
            # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
            transformer_norm_hidden_states = transformer_hidden_states
        elif self.transformer_block.use_ada_layer_norm_continuous:
            # transformer_norm_hidden_states = self.transformer_block.norm2(transformer_hidden_states, added_cond_kwargs["pooled_text_emb"])
            raise NotImplementedError()
        else:
            raise ValueError("Incorrect norm")
        if self.transformer_block.pos_embed is not None and self.transformer_block.use_ada_layer_norm_single is False:
            transformer_norm_hidden_states = self.transformer_block.pos_embed(transformer_norm_hidden_states)
        # ================================================== #
        # ================= CROSS ATTENTION ================ #
        # ================================================== #
        # We do an initial pass of the CrossAttention up to obtaining the similarity matrix here.
        # The similarity matrix is used to obtain scaling coefficients for the attention matrix of the self attention
        # We reuse the previously computed self-attention matrix, and only repeat the steps after the softmax
        cross_attention_input_hidden_states = (
            transformer_norm_hidden_states  # Renaming the variable for the sake of readability
        )
        # TODO: check if classifier_free_guidance is being used before splitting here
        if self.do_classifier_free_guidance:
            # Our scaling coefficients depend only on the conditional part, so we split the inputs
            (
                _cross_attention_input_hidden_states_unconditional,
                cross_attention_input_hidden_states_conditional,
            ) = cross_attention_input_hidden_states.chunk(2)
            # Same split for the encoder_hidden_states i.e. the tokens
            # Since the SelfAttention processors don't get the encoder states as input, we inject them into the processor in the begining.
            _encoder_hidden_states_unconditional, encoder_hidden_states_conditional = self.encoder_hidden_states.chunk(
                2
            )
        else:
            cross_attention_input_hidden_states_conditional = cross_attention_input_hidden_states
            encoder_hidden_states_conditional = self.encoder_hidden_states.chunk(2)
        # Rename the variables for the sake of readability
        # The part below is the beginning of the __call__ function of the following CrossAttention layer
        cross_attention_hidden_states = cross_attention_input_hidden_states_conditional
        cross_attention_encoder_hidden_states = encoder_hidden_states_conditional
        attn2 = self.transformer_block.attn2
        if attn2.spatial_norm is not None:
            cross_attention_hidden_states = attn2.spatial_norm(cross_attention_hidden_states, temb)
        input_ndim = cross_attention_hidden_states.ndim
        if input_ndim == 4:
            batch_size, channel, height, width = cross_attention_hidden_states.shape
            cross_attention_hidden_states = cross_attention_hidden_states.view(
                batch_size, channel, height * width
            ).transpose(1, 2)
        (
            batch_size,
            sequence_length,
            _,
        ) = cross_attention_hidden_states.shape  # It is definitely a cross attention, so no need for an if block
        # TODO: change the attention_mask here
        attention_mask = attn2.prepare_attention_mask(
            None, sequence_length, batch_size
        )  # I assume the attention mask is the same...
        if attn2.group_norm is not None:
            cross_attention_hidden_states = attn2.group_norm(cross_attention_hidden_states.transpose(1, 2)).transpose(
                1, 2
            )
        query2 = attn2.to_q(cross_attention_hidden_states)
        if attn2.norm_cross:
            cross_attention_encoder_hidden_states = attn2.norm_encoder_hidden_states(
                cross_attention_encoder_hidden_states
            )
        key2 = attn2.to_k(cross_attention_encoder_hidden_states)
        query2 = attn2.head_to_batch_dim(query2)
        key2 = attn2.head_to_batch_dim(key2)
        cross_attention_probs = attn2.get_attention_scores(query2, key2, attention_mask)
        # CrossAttention ends here, the remaining part is not used
        # ================================================== #
        # ================ SELF ATTENTION 2 ================ #
        # ================================================== #
        # DEJA VU!
        mask = (mask > 0.5).to(self_attention_output_hidden_states.dtype)
        m = mask.to(self_attention_output_hidden_states.device)
        # m = rearrange(m, 'b c h w -> b (h w) c').contiguous()
        m = m.permute(0, 2, 3, 1).reshape((m.shape[0], -1, m.shape[1])).contiguous()  # B HW 1
        m = torch.matmul(m, m.permute(0, 2, 1)) + (1 - m)
        # # Compute scaling coefficients for the similarity matrix
        # # Select the cross attention values for the correct tokens only!
        # cross_attention_probs = cross_attention_probs.mean(dim = 0)
        # cross_attention_probs = cross_attention_probs[:, self.token_idx].sum(dim=1)
        # cross_attention_probs = cross_attention_probs.reshape(shape)
        # gaussian_smoothing = GaussianSmoothing(channels=1, kernel_size=3, sigma=0.5, dim=2).to(self_attention_output_hidden_states.device)
        # cross_attention_probs = gaussian_smoothing(cross_attention_probs.unsqueeze(0))[0] # optional smoothing
        # cross_attention_probs = cross_attention_probs.reshape(-1)
        # cross_attention_probs = ((cross_attention_probs - torch.median(cross_attention_probs.ravel())) / torch.max(cross_attention_probs.ravel())).clip(0, 1)
        # c = (1 - m) * cross_attention_probs.reshape(1, 1, -1) + m # PAIntA scaling coefficients
        # Compute scaling coefficients for the similarity matrix
        # Select the cross attention values for the correct tokens only!
        batch_size, dims, channels = cross_attention_probs.shape
        batch_size = batch_size // attn.heads
        cross_attention_probs = cross_attention_probs.reshape((batch_size, attn.heads, dims, channels))  # B, D, HW, T
        cross_attention_probs = cross_attention_probs.mean(dim=1)  # B, HW, T
        cross_attention_probs = cross_attention_probs[..., self.token_idx].sum(dim=-1)  # B, HW
        cross_attention_probs = cross_attention_probs.reshape((batch_size,) + shape)  # , B, H, W
        gaussian_smoothing = GaussianSmoothing(channels=1, kernel_size=3, sigma=0.5, dim=2).to(
            self_attention_output_hidden_states.device
        )
        cross_attention_probs = gaussian_smoothing(cross_attention_probs[:, None])[:, 0]  # optional smoothing B, H, W
        # Median normalization
        cross_attention_probs = cross_attention_probs.reshape(batch_size, -1)  # B, HW
        cross_attention_probs = (
            cross_attention_probs - cross_attention_probs.median(dim=-1, keepdim=True).values
        ) / cross_attention_probs.max(dim=-1, keepdim=True).values
        cross_attention_probs = cross_attention_probs.clip(0, 1)
        c = (1 - m) * cross_attention_probs.reshape(batch_size, 1, -1) + m
        c = c.repeat_interleave(attn.heads, 0)  # BD, HW
        if self.do_classifier_free_guidance:
            c = torch.cat([c, c])  # 2BD, HW
        # Rescaling the original self-attention matrix
        self_attention_scores_rescaled = self_attention_scores * c
        self_attention_probs_rescaled = self_attention_scores_rescaled.softmax(dim=-1)
        # Continuing the self attention normally using the new matrix
        hidden_states = torch.bmm(self_attention_probs_rescaled, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)
        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)
        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
        if attn.residual_connection:
            hidden_states = hidden_states + input_hidden_states
        hidden_states = hidden_states / attn.rescale_output_factor
        return hidden_states
 class StableDiffusionHDPainterPipeline(StableDiffusionInpaintPipeline):
    def get_tokenized_prompt(self, prompt):
        out = self.tokenizer(prompt)
        return [self.tokenizer.decode(x) for x in out["input_ids"]]
    def init_attn_processors(
        self,
        mask,
        token_idx,
        use_painta=True,
        use_rasg=True,
        painta_scale_factors=[2, 4],  # 64x64 -> [16x16, 32x32]
        rasg_scale_factor=4,  # 64x64 -> 16x16
        self_attention_layer_name="attn1",
        cross_attention_layer_name="attn2",
        list_of_painta_layer_names=None,
        list_of_rasg_layer_names=None,
    ):
        default_processor = AttnProcessor()
        width, height = mask.shape[-2:]
        width, height = width // self.vae_scale_factor, height // self.vae_scale_factor
        painta_scale_factors = [x * self.vae_scale_factor for x in painta_scale_factors]
        rasg_scale_factor = self.vae_scale_factor * rasg_scale_factor
        attn_processors = {}
        for x in self.unet.attn_processors:
            if (list_of_painta_layer_names is None and self_attention_layer_name in x) or (
                list_of_painta_layer_names is not None and x in list_of_painta_layer_names
            ):
                if use_painta:
                    transformer_block = self.unet.get_submodule(x.replace(".attn1.processor", ""))
                    attn_processors[x] = PAIntAAttnProcessor(
                        transformer_block, mask, token_idx, self.do_classifier_free_guidance, painta_scale_factors
                    )
                else:
                    attn_processors[x] = default_processor
            elif (list_of_rasg_layer_names is None and cross_attention_layer_name in x) or (
                list_of_rasg_layer_names is not None and x in list_of_rasg_layer_names
            ):
                if use_rasg:
                    attn_processors[x] = RASGAttnProcessor(mask, token_idx, rasg_scale_factor)
                else:
                    attn_processors[x] = default_processor
        self.unet.set_attn_processor(attn_processors)
        # import json
        # with open('/home/hayk.manukyan/repos/diffusers/debug.txt', 'a')  as f:
        #     json.dump({x:str(y) for x,y in self.unet.attn_processors.items()}, f, indent=4)
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]] = None,
        image: PipelineImageInput = None,
        mask_image: PipelineImageInput = None,
        masked_image_latents: torch.FloatTensor = None,
        height: Optional[int] = None,
        width: Optional[int] = None,
        padding_mask_crop: Optional[int] = None,
        strength: float = 1.0,
        num_inference_steps: int = 50,
        timesteps: List[int] = None,
        guidance_scale: float = 7.5,
        positive_prompt: Optional[str] = "",
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_images_per_prompt: Optional[int] = 1,
        eta: float = 0.01,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.FloatTensor] = None,
        prompt_embeds: Optional[torch.FloatTensor] = None,
        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
        ip_adapter_image: Optional[PipelineImageInput] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
        clip_skip: int = None,
        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
        use_painta=True,
        use_rasg=True,
        self_attention_layer_name=".attn1",
        cross_attention_layer_name=".attn2",
        painta_scale_factors=[2, 4],  # 16 x 16 and 32 x 32
        rasg_scale_factor=4,  # 16x16 by default
        list_of_painta_layer_names=None,
        list_of_rasg_layer_names=None,
        **kwargs,
    ):
        callback = kwargs.pop("callback", None)
        callback_steps = kwargs.pop("callback_steps", None)
        if callback is not None:
            deprecate(
                "callback",
                "1.0.0",
                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
            )
        if callback_steps is not None:
            deprecate(
                "callback_steps",
                "1.0.0",
                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
            )
        # 0. Default height and width to unet
        height = height or self.unet.config.sample_size * self.vae_scale_factor
        width = width or self.unet.config.sample_size * self.vae_scale_factor
        #
        prompt_no_positives = prompt
        if isinstance(prompt, list):
            prompt = [x + positive_prompt for x in prompt]
        else:
            prompt = prompt + positive_prompt
        # 1. Check inputs
        self.check_inputs(
            prompt,
            image,
            mask_image,
            height,
            width,
            strength,
            callback_steps,
            negative_prompt,
            prompt_embeds,
            negative_prompt_embeds,
            callback_on_step_end_tensor_inputs,
            padding_mask_crop,
        )
        self._guidance_scale = guidance_scale
        self._clip_skip = clip_skip
        self._cross_attention_kwargs = cross_attention_kwargs
        self._interrupt = False
        # 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]
        # assert batch_size == 1, "Does not work with batch size > 1 currently"
        device = self._execution_device
        # 3. Encode input prompt
        text_encoder_lora_scale = (
            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
        )
        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            self.do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=text_encoder_lora_scale,
            clip_skip=self.clip_skip,
        )
        # For classifier free guidance, we need to do two forward passes.
        # Here we concatenate the unconditional and text embeddings into a single batch
        # to avoid doing two forward passes
        if self.do_classifier_free_guidance:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
        if ip_adapter_image is not None:
            output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True
            image_embeds, negative_image_embeds = self.encode_image(
                ip_adapter_image, device, num_images_per_prompt, output_hidden_state
            )
            if self.do_classifier_free_guidance:
                image_embeds = torch.cat([negative_image_embeds, image_embeds])
        # 4. set timesteps
        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
        timesteps, num_inference_steps = self.get_timesteps(
            num_inference_steps=num_inference_steps, strength=strength, device=device
        )
        # check that number of inference steps is not < 1 - as this doesn't make sense
        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."
            )
        # at which timestep to set the initial noise (n.b. 50% if strength is 0.5)
        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
        # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise
        is_strength_max = strength == 1.0
        # 5. Preprocess mask and image
        if padding_mask_crop is not None:
            crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop)
            resize_mode = "fill"
        else:
            crops_coords = None
            resize_mode = "default"
        original_image = image
        init_image = self.image_processor.preprocess(
            image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode
        )
        init_image = init_image.to(dtype=torch.float32)
        # 6. Prepare latent variables
        num_channels_latents = self.vae.config.latent_channels
        num_channels_unet = self.unet.config.in_channels
        return_image_latents = num_channels_unet == 4
        latents_outputs = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
            image=init_image,
            timestep=latent_timestep,
            is_strength_max=is_strength_max,
            return_noise=True,
            return_image_latents=return_image_latents,
        )
        if return_image_latents:
            latents, noise, image_latents = latents_outputs
        else:
            latents, noise = latents_outputs
        # 7. Prepare mask latent variables
        mask_condition = self.mask_processor.preprocess(
            mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords
        )
        if masked_image_latents is None:
            masked_image = init_image * (mask_condition < 0.5)
        else:
            masked_image = masked_image_latents
        mask, masked_image_latents = self.prepare_mask_latents(
            mask_condition,
            masked_image,
            batch_size * num_images_per_prompt,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            self.do_classifier_free_guidance,
        )
        # 7.5 Setting up HD-Painter
        # Get the indices of the tokens to be modified by both RASG and PAIntA
        token_idx = list(range(1, self.get_tokenized_prompt(prompt_no_positives).index("<|endoftext|>"))) + [
            self.get_tokenized_prompt(prompt).index("<|endoftext|>")
        ]
        # Setting up the attention processors
        self.init_attn_processors(
            mask_condition,
            token_idx,
            use_painta,
            use_rasg,
            painta_scale_factors=painta_scale_factors,
            rasg_scale_factor=rasg_scale_factor,
            self_attention_layer_name=self_attention_layer_name,
            cross_attention_layer_name=cross_attention_layer_name,
            list_of_painta_layer_names=list_of_painta_layer_names,
            list_of_rasg_layer_names=list_of_rasg_layer_names,
        )
        # 8. Check that sizes of mask, masked image and latents match
        if num_channels_unet == 9:
            # default case for runwayml/stable-diffusion-inpainting
            num_channels_mask = mask.shape[1]
            num_channels_masked_image = masked_image_latents.shape[1]
            if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels:
                raise ValueError(
                    f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects"
                    f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +"
                    f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}"
                    f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of"
                    " `pipeline.unet` or your `mask_image` or `image` input."
                )
        elif num_channels_unet != 4:
            raise ValueError(
                f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}."
            )
        # 9. 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)
        if use_rasg:
            extra_step_kwargs["generator"] = None
        # 9.1 Add image embeds for IP-Adapter
        added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None
        # 9.2 Optionally get Guidance Scale Embedding
        timestep_cond = None
        if self.unet.config.time_cond_proj_dim is not None:
            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
            timestep_cond = self.get_guidance_scale_embedding(
                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
            ).to(device=device, dtype=latents.dtype)
        # 10. Denoising loop
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        self._num_timesteps = len(timesteps)
        painta_active = True
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue
                if t < 500 and painta_active:
                    self.init_attn_processors(
                        mask_condition,
                        token_idx,
                        False,
                        use_rasg,
                        painta_scale_factors=painta_scale_factors,
                        rasg_scale_factor=rasg_scale_factor,
                        self_attention_layer_name=self_attention_layer_name,
                        cross_attention_layer_name=cross_attention_layer_name,
                        list_of_painta_layer_names=list_of_painta_layer_names,
                        list_of_rasg_layer_names=list_of_rasg_layer_names,
                    )
                    painta_active = False
                with torch.enable_grad():
                    self.unet.zero_grad()
                    latents = latents.detach()
                    latents.requires_grad = True
                    # expand the latents if we are doing classifier free guidance
                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                    # concat latents, mask, masked_image_latents in the channel dimension
                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
                    if num_channels_unet == 9:
                        latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1)
                    self.scheduler.latents = latents
                    self.encoder_hidden_states = prompt_embeds
                    for attn_processor in self.unet.attn_processors.values():
                        attn_processor.encoder_hidden_states = prompt_embeds
                    # predict the noise residual
                    noise_pred = self.unet(
                        latent_model_input,
                        t,
                        encoder_hidden_states=prompt_embeds,
                        timestep_cond=timestep_cond,
                        cross_attention_kwargs=self.cross_attention_kwargs,
                        added_cond_kwargs=added_cond_kwargs,
                        return_dict=False,
                    )[0]
                    # perform guidance
                    if self.do_classifier_free_guidance:
                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
                    if use_rasg:
                        # Perform RASG
                        _, _, height, width = mask_condition.shape  # 512 x 512
                        scale_factor = self.vae_scale_factor * rasg_scale_factor  # 8 * 4 = 32
                        # TODO: Fix for > 1 batch_size
                        rasg_mask = F.interpolate(
                            mask_condition, (height // scale_factor, width // scale_factor), mode="bicubic"
                        )[0, 0]  # mode is nearest by default, B, H, W
                        # Aggregate the saved attention maps
                        attn_map = []
                        for processor in self.unet.attn_processors.values():
                            if hasattr(processor, "attention_scores") and processor.attention_scores is not None:
                                if self.do_classifier_free_guidance:
                                    attn_map.append(processor.attention_scores.chunk(2)[1])  # (B/2) x H, 256, 77
                                else:
                                    attn_map.append(processor.attention_scores)  # B x H, 256, 77 ?
                        attn_map = (
                            torch.cat(attn_map)
                            .mean(0)
                            .permute(1, 0)
                            .reshape((-1, height // scale_factor, width // scale_factor))
                        )  # 77, 16, 16
                        # Compute the attention score
                        attn_score = -sum(
                            [
                                F.binary_cross_entropy_with_logits(x - 1.0, rasg_mask.to(device))
                                for x in attn_map[token_idx]
                            ]
                        )
                        # Backward the score and compute the gradients
                        attn_score.backward()
                        # Normalzie the gradients and compute the noise component
                        variance_noise = latents.grad.detach()
                        # print("VARIANCE SHAPE", variance_noise.shape)
                        variance_noise -= torch.mean(variance_noise, [1, 2, 3], keepdim=True)
                        variance_noise /= torch.std(variance_noise, [1, 2, 3], keepdim=True)
                    else:
                        variance_noise = None
                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(
                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False, variance_noise=variance_noise
                )[0]
                if num_channels_unet == 4:
                    init_latents_proper = image_latents
                    if self.do_classifier_free_guidance:
                        init_mask, _ = mask.chunk(2)
                    else:
                        init_mask = mask
                    if i < len(timesteps) - 1:
                        noise_timestep = timesteps[i + 1]
                        init_latents_proper = self.scheduler.add_noise(
                            init_latents_proper, noise, torch.tensor([noise_timestep])
                        )
                    latents = (1 - init_mask) * init_latents_proper + init_mask * latents
                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)
                    mask = callback_outputs.pop("mask", mask)
                    masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents)
                # 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 callback is not None and i % callback_steps == 0:
                        step_idx = i // getattr(self.scheduler, "order", 1)
                        callback(step_idx, t, latents)
        if not output_type == "latent":
            condition_kwargs = {}
            if isinstance(self.vae, AsymmetricAutoencoderKL):
                init_image = init_image.to(device=device, dtype=masked_image_latents.dtype)
                init_image_condition = init_image.clone()
                init_image = self._encode_vae_image(init_image, generator=generator)
                mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype)
                condition_kwargs = {"image": init_image_condition, "mask": mask_condition}
            image = self.vae.decode(
                latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs
            )[0]
            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
        else:
            image = latents
            has_nsfw_concept = None
        if has_nsfw_concept is None:
            do_denormalize = [True] * image.shape[0]
        else:
            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
        if padding_mask_crop is not None:
            image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image]
        # Offload all models
        self.maybe_free_model_hooks()
        if not return_dict:
            return (image, has_nsfw_concept)
        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
 # ============= Utility Functions ============== #
 class GaussianSmoothing(nn.Module):
    """
    Apply gaussian smoothing on a
    1d, 2d or 3d tensor. Filtering is performed seperately for each channel
    in the input using a depthwise convolution.
    Arguments:
        channels (int, sequence): Number of channels of the input tensors. Output will
            have this number of channels as well.
        kernel_size (int, sequence): Size of the gaussian kernel.
        sigma (float, sequence): Standard deviation of the gaussian kernel.
        dim (int, optional): The number of dimensions of the data.
            Default value is 2 (spatial).
    """
    def __init__(self, channels, kernel_size, sigma, dim=2):
        super(GaussianSmoothing, self).__init__()
        if isinstance(kernel_size, numbers.Number):
            kernel_size = [kernel_size] * dim
        if isinstance(sigma, numbers.Number):
            sigma = [sigma] * dim
        # The gaussian kernel is the product of the
        # gaussian function of each dimension.
        kernel = 1
        meshgrids = torch.meshgrid([torch.arange(size, dtype=torch.float32) for size in kernel_size])
        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
            mean = (size - 1) / 2
            kernel *= 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-(((mgrid - mean) / (2 * std)) ** 2))
        # Make sure sum of values in gaussian kernel equals 1.
        kernel = kernel / torch.sum(kernel)
        # Reshape to depthwise convolutional weight
        kernel = kernel.view(1, 1, *kernel.size())
        kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
        self.register_buffer("weight", kernel)
        self.groups = channels
        if dim == 1:
            self.conv = F.conv1d
        elif dim == 2:
            self.conv = F.conv2d
        elif dim == 3:
            self.conv = F.conv3d
        else:
            raise RuntimeError("Only 1, 2 and 3 dimensions are supported. Received {}.".format(dim))
    def forward(self, input):
        """
        Apply gaussian filter to input.
        Arguments:
            input (torch.Tensor): Input to apply gaussian filter on.
        Returns:
            filtered (torch.Tensor): Filtered output.
        """
        return self.conv(input, weight=self.weight.to(input.dtype), groups=self.groups, padding="same")
 def get_attention_scores(
    self, query: torch.Tensor, key: torch.Tensor, attention_mask: torch.Tensor = None
 ) -> torch.Tensor:
    r"""
    Compute the attention scores.
    Args:
        query (`torch.Tensor`): The query tensor.
        key (`torch.Tensor`): The key tensor.
        attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
    Returns:
        `torch.Tensor`: The attention probabilities/scores.
    """
    if self.upcast_attention:
        query = query.float()
        key = key.float()
    if attention_mask is None:
        baddbmm_input = torch.empty(
            query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
        )
        beta = 0
    else:
        baddbmm_input = attention_mask
        beta = 1
    attention_scores = torch.baddbmm(
        baddbmm_input,
        query,
        key.transpose(-1, -2),
        beta=beta,
        alpha=self.scale,
    )
    del baddbmm_input
    if self.upcast_softmax:
        attention_scores = attention_scores.float()
    return attention_scores
@@ -1,8 +1,7 @@
 """
-modeled after the textual_inversion.py / train_dreambooth.py and the work
+    modeled after the textual_inversion.py / train_dreambooth.py and the work
-of justinpinkney here: https://github.com/justinpinkney/stable-diffusion/blob/main/notebooks/imagic.ipynb
+    of justinpinkney here: https://github.com/justinpinkney/stable-diffusion/blob/main/notebooks/imagic.ipynb
 """
 import inspect
 import warnings
 from typing import List, Optional, Union
@@ -440,7 +440,7 @@ def betas_for_alpha_bar(
            return math.exp(t * -12.0)
    else:
-        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
    betas = []
    for i in range(num_diffusion_timesteps):
@@ -726,7 +726,7 @@ class LatentConsistencyModelWalkPipeline(
            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.
+                `._callback_tensor_inputs` attribute of your pipeine class.
            embedding_interpolation_type (`str`, *optional*, defaults to `"lerp"`):
                The type of interpolation to use for interpolating between text embeddings. Choose between `"lerp"` and `"slerp"`.
            latent_interpolation_type (`str`, *optional*, defaults to `"slerp"`):
@@ -779,7 +779,7 @@ class LatentConsistencyModelWalkPipeline(
        else:
            batch_size = prompt_embeds.shape[0]
        if batch_size < 2:
-            raise ValueError(f"`prompt` must have length of at least 2 but found {batch_size}")
+            raise ValueError(f"`prompt` must have length of atleast 2 but found {batch_size}")
        if num_images_per_prompt != 1:
            raise ValueError("`num_images_per_prompt` must be `1` as no other value is supported yet")
        if prompt_embeds is not None:
@@ -883,7 +883,7 @@ class LatentConsistencyModelWalkPipeline(
                ) as batch_progress_bar:
                    for batch_index in range(0, bs, process_batch_size):
                        batch_inference_latents = inference_latents[batch_index : batch_index + process_batch_size]
-                        batch_inference_embeddings = inference_embeddings[
+                        batch_inference_embedddings = inference_embeddings[
                            batch_index : batch_index + process_batch_size
                        ]
@@ -892,7 +892,7 @@ class LatentConsistencyModelWalkPipeline(
                        )
                        timesteps = self.scheduler.timesteps
-                        current_bs = batch_inference_embeddings.shape[0]
+                        current_bs = batch_inference_embedddings.shape[0]
                        w = torch.tensor(self.guidance_scale - 1).repeat(current_bs)
                        w_embedding = self.get_guidance_scale_embedding(
                            w, embedding_dim=self.unet.config.time_cond_proj_dim
@@ -901,14 +901,14 @@ class LatentConsistencyModelWalkPipeline(
                        # 10. Perform inference for current batch
                        with self.progress_bar(total=num_inference_steps) as progress_bar:
                            for index, t in enumerate(timesteps):
-                                batch_inference_latents = batch_inference_latents.to(batch_inference_embeddings.dtype)
+                                batch_inference_latents = batch_inference_latents.to(batch_inference_embedddings.dtype)
                                # model prediction (v-prediction, eps, x)
                                model_pred = self.unet(
                                    batch_inference_latents,
                                    t,
                                    timestep_cond=w_embedding,
-                                    encoder_hidden_states=batch_inference_embeddings,
+                                    encoder_hidden_states=batch_inference_embedddings,
                                    cross_attention_kwargs=self.cross_attention_kwargs,
                                    return_dict=False,
                                )[0]
@@ -924,8 +924,8 @@ class LatentConsistencyModelWalkPipeline(
                                    callback_outputs = callback_on_step_end(self, index, t, callback_kwargs)
                                    batch_inference_latents = callback_outputs.pop("latents", batch_inference_latents)
-                                    batch_inference_embeddings = callback_outputs.pop(
+                                    batch_inference_embedddings = callback_outputs.pop(
-                                        "prompt_embeds", batch_inference_embeddings
+                                        "prompt_embeds", batch_inference_embedddings
                                    )
                                    w_embedding = callback_outputs.pop("w_embedding", w_embedding)
                                    denoised = callback_outputs.pop("denoised", denoised)
@@ -939,7 +939,7 @@ class LatentConsistencyModelWalkPipeline(
                                        step_idx = index // getattr(self.scheduler, "order", 1)
                                        callback(step_idx, t, batch_inference_latents)
-                        denoised = denoised.to(batch_inference_embeddings.dtype)
+                        denoised = denoised.to(batch_inference_embedddings.dtype)
                        # Note: This is not supported because you would get black images in your latent walk if
                        #       NSFW concept is detected
@@ -348,7 +348,7 @@ def betas_for_alpha_bar(
            return math.exp(t * -12.0)
    else:
-        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
    betas = []
    for i in range(num_diffusion_timesteps):
@@ -530,7 +530,7 @@ class LLMGroundedDiffusionPipeline(
                )
        if len(phrases) != len(boxes):
-            raise ValueError(
+            ValueError(
                "length of `phrases` and `boxes` has to be same, but"
                f" got: `phrases` {len(phrases)} != `boxes` {len(boxes)}"
            )
@@ -164,7 +164,7 @@ def get_prompts_tokens_with_weights(clip_tokenizer: CLIPTokenizer, prompt: str):
        text_tokens (list)
            A list contains token ids
        text_weight (list)
-            A list contains the correspondent weight of token ids
+            A list contains the correspodent weight of token ids
    Example:
        import torch
@@ -1028,7 +1028,7 @@ class SDXLLongPromptWeightingPipeline(
                # because `num_inference_steps` might be even given that every timestep
                # (except the highest one) is duplicated. If `num_inference_steps` is even it would
                # mean that we cut the timesteps in the middle of the denoising step
-                # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1
+                # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
                # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
                num_inference_steps = num_inference_steps + 1
@@ -1531,7 +1531,7 @@ class SDXLLongPromptWeightingPipeline(
            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.
+                `._callback_tensor_inputs` attribute of your pipeine class.
        Examples:
@@ -2131,7 +2131,7 @@ class SDXLLongPromptWeightingPipeline(
            **kwargs,
        )
-    # Override to properly handle the loading and unloading of the additional text encoder.
+    # Overrride to properly handle the loading and unloading of the additional text encoder.
    def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
        # We could have accessed the unet config from `lora_state_dict()` too. We pass
        # it here explicitly to be able to tell that it's coming from an SDXL
@@ -18,7 +18,6 @@
 # --------------------------------------------------------------------------
 import logging
 import math
 from typing import Dict, Union
@@ -26,7 +25,6 @@ import matplotlib
 import numpy as np
 import torch
 from PIL import Image
 from PIL.Image import Resampling
 from scipy.optimize import minimize
 from torch.utils.data import DataLoader, TensorDataset
 from tqdm.auto import tqdm
@@ -36,14 +34,13 @@ from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LCMScheduler,
    UNet2DConditionModel,
 )
 from diffusers.utils import BaseOutput, check_min_version
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.25.0")
+check_min_version("0.27.0.dev0")
 class MarigoldDepthOutput(BaseOutput):
@@ -64,19 +61,6 @@ class MarigoldDepthOutput(BaseOutput):
    uncertainty: Union[None, np.ndarray]
 def get_pil_resample_method(method_str: str) -> Resampling:
    resample_method_dic = {
        "bilinear": Resampling.BILINEAR,
        "bicubic": Resampling.BICUBIC,
        "nearest": Resampling.NEAREST,
    }
    resample_method = resample_method_dic.get(method_str, None)
    if resample_method is None:
        raise ValueError(f"Unknown resampling method: {resample_method}")
    else:
        return resample_method
 class MarigoldPipeline(DiffusionPipeline):
    """
    Pipeline for monocular depth estimation using Marigold: https://marigoldmonodepth.github.io.
@@ -129,9 +113,7 @@ class MarigoldPipeline(DiffusionPipeline):
        ensemble_size: int = 10,
        processing_res: int = 768,
        match_input_res: bool = True,
        resample_method: str = "bilinear",
        batch_size: int = 0,
        seed: Union[int, None] = None,
        color_map: str = "Spectral",
        show_progress_bar: bool = True,
        ensemble_kwargs: Dict = None,
@@ -147,9 +129,7 @@ class MarigoldPipeline(DiffusionPipeline):
                If set to 0: will not resize at all.
            match_input_res (`bool`, *optional*, defaults to `True`):
                Resize depth prediction to match input resolution.
-                Only valid if `processing_res` > 0.
+                Only valid if `limit_input_res` is not None.
            resample_method: (`str`, *optional*, defaults to `bilinear`):
                Resampling method used to resize images and depth predictions. This can be one of `bilinear`, `bicubic` or `nearest`, defaults to: `bilinear`.
            denoising_steps (`int`, *optional*, defaults to `10`):
                Number of diffusion denoising steps (DDIM) during inference.
            ensemble_size (`int`, *optional*, defaults to `10`):
@@ -157,8 +137,6 @@ class MarigoldPipeline(DiffusionPipeline):
            batch_size (`int`, *optional*, defaults to `0`):
                Inference batch size, no bigger than `num_ensemble`.
                If set to 0, the script will automatically decide the proper batch size.
            seed (`int`, *optional*, defaults to `None`)
                Reproducibility seed.
            show_progress_bar (`bool`, *optional*, defaults to `True`):
                Display a progress bar of diffusion denoising.
            color_map (`str`, *optional*, defaults to `"Spectral"`, pass `None` to skip colorized depth map generation):
@@ -168,7 +146,8 @@ class MarigoldPipeline(DiffusionPipeline):
        Returns:
            `MarigoldDepthOutput`: Output class for Marigold monocular depth prediction pipeline, including:
            - **depth_np** (`np.ndarray`) Predicted depth map, with depth values in the range of [0, 1]
-            - **depth_colored** (`PIL.Image.Image`) Colorized depth map, with the shape of [3, H, W] and values in [0, 1], None if `color_map` is `None`
+            - **depth_colored** (`None` or `PIL.Image.Image`) Colorized depth map, with the shape of [3, H, W] and
                    values in [0, 1]. None if `color_map` is `None`
            - **uncertainty** (`None` or `np.ndarray`) Uncalibrated uncertainty(MAD, median absolute deviation)
                    coming from ensembling. None if `ensemble_size = 1`
        """
@@ -179,21 +158,13 @@ class MarigoldPipeline(DiffusionPipeline):
        if not match_input_res:
            assert processing_res is not None, "Value error: `resize_output_back` is only valid with "
        assert processing_res >= 0
        assert denoising_steps >= 1
        assert ensemble_size >= 1
        # Check if denoising step is reasonable
        self._check_inference_step(denoising_steps)
        resample_method: Resampling = get_pil_resample_method(resample_method)
        # ----------------- Image Preprocess -----------------
        # Resize image
        if processing_res > 0:
-            input_image = self.resize_max_res(
+            input_image = self.resize_max_res(input_image, max_edge_resolution=processing_res)
                input_image,
                max_edge_resolution=processing_res,
                resample_method=resample_method,
            )
        # Convert the image to RGB, to 1.remove the alpha channel 2.convert B&W to 3-channel
        input_image = input_image.convert("RGB")
        image = np.asarray(input_image)
@@ -232,10 +203,9 @@ class MarigoldPipeline(DiffusionPipeline):
                rgb_in=batched_img,
                num_inference_steps=denoising_steps,
                show_pbar=show_progress_bar,
                seed=seed,
            )
-            depth_pred_ls.append(depth_pred_raw.detach())
+            depth_pred_ls.append(depth_pred_raw.detach().clone())
-        depth_preds = torch.concat(depth_pred_ls, dim=0).squeeze()
+        depth_preds = torch.concat(depth_pred_ls, axis=0).squeeze()
        torch.cuda.empty_cache()  # clear vram cache for ensembling
        # ----------------- Test-time ensembling -----------------
@@ -257,7 +227,7 @@ class MarigoldPipeline(DiffusionPipeline):
        # Resize back to original resolution
        if match_input_res:
            pred_img = Image.fromarray(depth_pred)
-            pred_img = pred_img.resize(input_size, resample=resample_method)
+            pred_img = pred_img.resize(input_size)
            depth_pred = np.asarray(pred_img)
        # Clip output range
@@ -273,32 +243,12 @@ class MarigoldPipeline(DiffusionPipeline):
            depth_colored_img = Image.fromarray(depth_colored_hwc)
        else:
            depth_colored_img = None
        return MarigoldDepthOutput(
            depth_np=depth_pred,
            depth_colored=depth_colored_img,
            uncertainty=pred_uncert,
        )
    def _check_inference_step(self, n_step: int):
        """
        Check if denoising step is reasonable
        Args:
            n_step (`int`): denoising steps
        """
        assert n_step >= 1
        if isinstance(self.scheduler, DDIMScheduler):
            if n_step < 10:
                logging.warning(
                    f"Too few denoising steps: {n_step}. Recommended to use the LCM checkpoint for few-step inference."
                )
        elif isinstance(self.scheduler, LCMScheduler):
            if not 1 <= n_step <= 4:
                logging.warning(f"Non-optimal setting of denoising steps: {n_step}. Recommended setting is 1-4 steps.")
        else:
            raise RuntimeError(f"Unsupported scheduler type: {type(self.scheduler)}")
    def _encode_empty_text(self):
        """
        Encode text embedding for empty prompt.
@@ -315,13 +265,7 @@ class MarigoldPipeline(DiffusionPipeline):
        self.empty_text_embed = self.text_encoder(text_input_ids)[0].to(self.dtype)
    @torch.no_grad()
-    def single_infer(
+    def single_infer(self, rgb_in: torch.Tensor, num_inference_steps: int, show_pbar: bool) -> torch.Tensor:
        self,
        rgb_in: torch.Tensor,
        num_inference_steps: int,
        seed: Union[int, None],
        show_pbar: bool,
    ) -> torch.Tensor:
        """
        Perform an individual depth prediction without ensembling.
@@ -342,20 +286,10 @@ class MarigoldPipeline(DiffusionPipeline):
        timesteps = self.scheduler.timesteps  # [T]
        # Encode image
-        rgb_latent = self.encode_rgb(rgb_in)
+        rgb_latent = self._encode_rgb(rgb_in)
        # Initial depth map (noise)
-        if seed is None:
+        depth_latent = torch.randn(rgb_latent.shape, device=device, dtype=self.dtype)  # [B, 4, h, w]
            rand_num_generator = None
        else:
            rand_num_generator = torch.Generator(device=device)
            rand_num_generator.manual_seed(seed)
        depth_latent = torch.randn(
            rgb_latent.shape,
            device=device,
            dtype=self.dtype,
            generator=rand_num_generator,
        )  # [B, 4, h, w]
        # Batched empty text embedding
        if self.empty_text_embed is None:
@@ -380,9 +314,9 @@ class MarigoldPipeline(DiffusionPipeline):
            noise_pred = self.unet(unet_input, t, encoder_hidden_states=batch_empty_text_embed).sample  # [B, 4, h, w]
            # compute the previous noisy sample x_t -> x_t-1
-            depth_latent = self.scheduler.step(noise_pred, t, depth_latent, generator=rand_num_generator).prev_sample
+            depth_latent = self.scheduler.step(noise_pred, t, depth_latent).prev_sample
-
+        torch.cuda.empty_cache()
-        depth = self.decode_depth(depth_latent)
+        depth = self._decode_depth(depth_latent)
        # clip prediction
        depth = torch.clip(depth, -1.0, 1.0)
@@ -391,7 +325,7 @@ class MarigoldPipeline(DiffusionPipeline):
        return depth
-    def encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
+    def _encode_rgb(self, rgb_in: torch.Tensor) -> torch.Tensor:
        """
        Encode RGB image into latent.
@@ -410,7 +344,7 @@ class MarigoldPipeline(DiffusionPipeline):
        rgb_latent = mean * self.rgb_latent_scale_factor
        return rgb_latent
-    def decode_depth(self, depth_latent: torch.Tensor) -> torch.Tensor:
+    def _decode_depth(self, depth_latent: torch.Tensor) -> torch.Tensor:
        """
        Decode depth latent into depth map.
@@ -431,7 +365,7 @@ class MarigoldPipeline(DiffusionPipeline):
        return depth_mean
    @staticmethod
-    def resize_max_res(img: Image.Image, max_edge_resolution: int, resample_method=Resampling.BILINEAR) -> Image.Image:
+    def resize_max_res(img: Image.Image, max_edge_resolution: int) -> Image.Image:
        """
        Resize image to limit maximum edge length while keeping aspect ratio.
@@ -440,8 +374,6 @@ class MarigoldPipeline(DiffusionPipeline):
                Image to be resized.
            max_edge_resolution (`int`):
                Maximum edge length (pixel).
            resample_method (`PIL.Image.Resampling`):
                Resampling method used to resize images.
        Returns:
            `Image.Image`: Resized image.
@@ -452,7 +384,7 @@ class MarigoldPipeline(DiffusionPipeline):
        new_width = int(original_width * downscale_factor)
        new_height = int(original_height * downscale_factor)
-        resized_img = img.resize((new_width, new_height), resample=resample_method)
+        resized_img = img.resize((new_width, new_height))
        return resized_img
    @staticmethod
@@ -196,7 +196,7 @@ class StableDiffusionTilingPipeline(DiffusionPipeline, StableDiffusionExtrasMixi
            guidance_scale_tiles: specific weights for classifier-free guidance in each tile.
            guidance_scale_tiles: specific weights for classifier-free guidance in each tile. If None, the value provided in guidance_scale will be used.
            seed_tiles: specific seeds for the initialization latents in each tile. These will override the latents generated for the whole canvas using the standard seed parameter.
-            seed_tiles_mode: either "full" "exclusive". If "full", all the latents affected by the tile be overriden. If "exclusive", only the latents that are affected exclusively by this tile (and no other tiles) will be overriden.
+            seed_tiles_mode: either "full" "exclusive". If "full", all the latents affected by the tile be overriden. If "exclusive", only the latents that are affected exclusively by this tile (and no other tiles) will be overrriden.
            seed_reroll_regions: a list of tuples in the form (start row, end row, start column, end column, seed) defining regions in pixel space for which the latents will be overriden using the given seed. Takes priority over seed_tiles.
            cpu_vae: the decoder from latent space to pixel space can require too mucho GPU RAM for large images. If you find out of memory errors at the end of the generation process, try setting this parameter to True to run the decoder in CPU. Slower, but should run without memory issues.
@@ -325,7 +325,7 @@ class StableDiffusionTilingPipeline(DiffusionPipeline, StableDiffusionExtrasMixi
        if accepts_eta:
            extra_step_kwargs["eta"] = eta
-        # Mask for tile weights strength
+        # Mask for tile weights strenght
        tile_weights = self._gaussian_weights(tile_width, tile_height, batch_size)
        # Diffusion timesteps
@@ -832,7 +832,7 @@ class AnimateDiffControlNetPipeline(
            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`, *optional*):
+            allback_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
@@ -840,7 +840,7 @@ class AnimateDiffControlNetPipeline(
            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.
+                `._callback_tensor_inputs` attribute of your pipeine class.
        Examples:
@@ -1280,7 +1280,7 @@ class DemoFusionSDXLPipeline(
        return output_images
-    # Override to properly handle the loading and unloading of the additional text encoder.
+    # Overrride to properly handle the loading and unloading of the additional text encoder.
    def load_lora_weights(self, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], **kwargs):
        # We could have accessed the unet config from `lora_state_dict()` too. We pass
        # it here explicitly to be able to tell that it's coming from an SDXL
@@ -887,7 +887,7 @@ class StyleAlignedSDXLPipeline(
                # because `num_inference_steps` might be even given that every timestep
                # (except the highest one) is duplicated. If `num_inference_steps` is even it would
                # mean that we cut the timesteps in the middle of the denoising step
-                # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1
+                # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
                # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
                num_inference_steps = num_inference_steps + 1
@@ -26,7 +26,7 @@ from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInver
 from diffusers.models import AutoencoderKL, UNet2DConditionModel
 from diffusers.models.lora import adjust_lora_scale_text_encoder
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
-from diffusers.pipelines.stable_diffusion_ldm3d.pipeline_stable_diffusion_ldm3d import LDM3DPipelineOutput
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_ldm3d import LDM3DPipelineOutput
 from diffusers.schedulers import DDPMScheduler, KarrasDiffusionSchedulers
 from diffusers.utils import (
    USE_PEFT_BACKEND,
@@ -206,7 +206,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool
            dimensions: ``batch x channels x height x width``.
    """
-    # checkpoint. #TODO(Yiyi) - need to clean this up later
+    # checkpoint. TOD(Yiyi) - need to clean this up later
    if image is None:
        raise ValueError("`image` input cannot be undefined.")
@@ -277,7 +277,7 @@ def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool
        # images are in latent space and thus can't
        # be masked set masked_image to None
        # we assume that the checkpoint is not an inpainting
-        # checkpoint. #TODO(Yiyi) - need to clean this up later
+        # checkpoint. TOD(Yiyi) - need to clean this up later
        masked_image = None
    else:
        masked_image = image * (mask < 0.5)
@@ -1073,7 +1073,7 @@ class StableDiffusionXLControlNetAdapterInpaintPipeline(
                # because `num_inference_steps` might be even given that every timestep
                # (except the highest one) is duplicated. If `num_inference_steps` is even it would
                # mean that we cut the timesteps in the middle of the denoising step
-                # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1
+                # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1
                # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
                num_inference_steps = num_inference_steps + 1
@@ -46,11 +46,6 @@ except Exception:
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 logger.warning(
    "To use instant id pipelines, please make sure you have the `insightface` library installed: `pip install insightface`."
    "Please refer to: https://huggingface.co/InstantX/InstantID for further instructions regarding inference"
 )
 def FeedForward(dim, mult=4):
    inner_dim = int(dim * mult)
@@ -706,7 +701,7 @@ class StableDiffusionXLInstantIDPipeline(StableDiffusionXLControlNetPipeline):
            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.
+                `._callback_tensor_inputs` attribute of your pipeine class.
        Examples:
@@ -81,7 +81,7 @@ def betas_for_alpha_bar(
            return math.exp(t * -12.0)
    else:
-        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
    betas = []
    for i in range(num_diffusion_timesteps):
@@ -1,7 +1,6 @@
 """
-modified based on diffusion library from Huggingface: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+    modified based on diffusion library from Huggingface: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
 """
 import inspect
 from typing import Callable, List, Optional, Union
@@ -224,7 +224,7 @@ class StableDiffusionIPEXPipeline(
        # 5. Prepare latent variables
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
-            self.unet.config.in_channels,
+            self.unet.in_channels,
            height,
            width,
            prompt_embeds.dtype,
@@ -679,7 +679,7 @@ class StableDiffusionIPEXPipeline(
        timesteps = self.scheduler.timesteps
        # 5. Prepare latent variables
-        num_channels_latents = self.unet.config.in_channels
+        num_channels_latents = self.unet.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
@@ -917,7 +917,7 @@ class TensorRTStableDiffusionPipeline(StableDiffusionPipeline):
            text_embeddings = self.__encode_prompt(prompt, negative_prompt)
            # Pre-initialize latents
-            num_channels_latents = self.unet.config.in_channels
+            num_channels_latents = self.unet.in_channels
            latents = self.prepare_latents(
                batch_size,
                num_channels_latents,
@@ -35,6 +35,7 @@ def slerp(val, low, high):
 class UnCLIPTextInterpolationPipeline(DiffusionPipeline):
    """
    Pipeline for prompt-to-prompt interpolation on CLIP text embeddings and using the UnCLIP / Dall-E to decode them to images.
@@ -48,7 +49,7 @@ class UnCLIPTextInterpolationPipeline(DiffusionPipeline):
            Tokenizer of class
            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
        prior ([`PriorTransformer`]):
-            The canonical unCLIP prior to approximate the image embedding from the text embedding.
+            The canonincal unCLIP prior to approximate the image embedding from the text embedding.
        text_proj ([`UnCLIPTextProjModel`]):
            Utility class to prepare and combine the embeddings before they are passed to the decoder.
        decoder ([`UNet2DConditionModel`]):
@@ -72,7 +72,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -65,7 +65,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -78,7 +78,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -71,7 +71,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -77,7 +77,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -60,7 +60,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -60,7 +60,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = logging.getLogger(__name__)
@@ -61,7 +61,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -125,11 +125,7 @@ def log_validation(vae, unet, controlnet, args, accelerator, weight_dtype, step,
        )
    image_logs = []
-    inference_ctx = (
+    inference_ctx = contextlib.nullcontext() if is_final_validation else torch.autocast("cuda")
        contextlib.nullcontext()
        if (is_final_validation or torch.backends.mps.is_available())
        else torch.autocast("cuda")
    )
    for validation_prompt, validation_image in zip(validation_prompts, validation_images):
        validation_image = Image.open(validation_image).convert("RGB")
@@ -796,12 +792,6 @@ def main(args):
    logging_dir = Path(args.output_dir, args.logging_dir)
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    accelerator = Accelerator(
@@ -63,7 +63,7 @@ from diffusers.utils.import_utils import is_xformers_available
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -259,9 +259,9 @@ The authors found that by using DoRA, both the learning capacity and training st
 > This is also aligned with some of the quantitative analysis shown in the paper. 
 **Usage**
-1. To use DoRA you need to upgrade the installation of `peft`: 
+1. To use DoRA you need to install `peft` from main: 
 ```bash
-pip install-U peft
+pip install git+https://github.com/huggingface/peft.git
 ```
 2. Enable DoRA training by adding this flag
 ```bash
@@ -269,7 +269,3 @@ pip install-U peft
 ```
 **Inference** 
 The inference is the same as if you train a regular LoRA 🤗
 ## Format compatibility
 You can pass `--output_kohya_format` to additionally generate a state dictionary which should be compatible with other platforms and tools such as Automatic 1111, Comfy, Kohya, etc. The `output_dir` will contain a file named "pytorch_lora_weights_kohya.safetensors".
@@ -63,7 +63,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -35,7 +35,7 @@ from diffusers.utils import check_min_version
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 # Cache compiled models across invocations of this script.
 cc.initialize_cache(os.path.expanduser("~/.cache/jax/compilation_cache"))
@@ -70,7 +70,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 import argparse
 import contextlib
 import gc
 import itertools
 import json
@@ -40,7 +41,6 @@ from peft import LoraConfig, set_peft_model_state_dict
 from peft.utils import get_peft_model_state_dict
 from PIL import Image
 from PIL.ImageOps import exif_transpose
 from safetensors.torch import load_file, save_file
 from torch.utils.data import Dataset
 from torchvision import transforms
 from torchvision.transforms.functional import crop
@@ -62,9 +62,7 @@ from diffusers.optimization import get_scheduler
 from diffusers.training_utils import _set_state_dict_into_text_encoder, cast_training_params, compute_snr
 from diffusers.utils import (
    check_min_version,
    convert_all_state_dict_to_peft,
    convert_state_dict_to_diffusers,
    convert_state_dict_to_kohya,
    convert_unet_state_dict_to_peft,
    is_wandb_available,
 )
@@ -77,7 +75,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -207,18 +205,11 @@ def log_validation(
    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
    # Currently the context determination is a bit hand-wavy. We can improve it in the future if there's a better
    # way to condition it. Reference: https://github.com/huggingface/diffusers/pull/7126#issuecomment-1968523051
-    enable_autocast = True
+    inference_ctx = (
-    if torch.backends.mps.is_available() or (
+        contextlib.nullcontext() if "playground" in args.pretrained_model_name_or_path else torch.cuda.amp.autocast()
-        accelerator.mixed_precision == "fp16" or accelerator.mixed_precision == "bf16"
+    )
    ):
        enable_autocast = False
    if "playground" in args.pretrained_model_name_or_path:
        enable_autocast = False
-    with torch.autocast(
+    with inference_ctx:
        accelerator.device.type,
        enabled=enable_autocast,
    ):
        images = [pipeline(**pipeline_args, generator=generator).images[0] for _ in range(args.num_validation_images)]
    for tracker in accelerator.trackers:
@@ -236,8 +227,7 @@ def log_validation(
            )
    del pipeline
-    if torch.cuda.is_available():
+    torch.cuda.empty_cache()
        torch.cuda.empty_cache()
    return images
@@ -406,11 +396,6 @@ def parse_args(input_args=None):
        default="lora-dreambooth-model",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--output_kohya_format",
        action="store_true",
        help="Flag to additionally generate final state dict in the Kohya format so that it becomes compatible with A111, Comfy, Kohya, etc.",
    )
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument(
        "--resolution",
@@ -974,12 +959,6 @@ def main(args):
    if args.do_edm_style_training and args.snr_gamma is not None:
        raise ValueError("Min-SNR formulation is not supported when conducting EDM-style training.")
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    logging_dir = Path(args.output_dir, args.logging_dir)
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
@@ -1022,8 +1001,7 @@ def main(args):
        cur_class_images = len(list(class_images_dir.iterdir()))
        if cur_class_images < args.num_class_images:
-            has_supported_fp16_accelerator = torch.cuda.is_available() or torch.backends.mps.is_available()
+            torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32
            torch_dtype = torch.float16 if has_supported_fp16_accelerator else torch.float32
            if args.prior_generation_precision == "fp32":
                torch_dtype = torch.float32
            elif args.prior_generation_precision == "fp16":
@@ -1148,12 +1126,6 @@ def main(args):
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
    if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    # Move unet, vae and text_encoder to device and cast to weight_dtype
    unet.to(accelerator.device, dtype=weight_dtype)
@@ -1298,7 +1270,7 @@ def main(args):
    # Enable TF32 for faster training on Ampere GPUs,
    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
-    if args.allow_tf32 and torch.cuda.is_available():
+    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True
    if args.scale_lr:
@@ -1475,8 +1447,7 @@ def main(args):
    if not args.train_text_encoder and not train_dataset.custom_instance_prompts:
        del tokenizers, text_encoders
        gc.collect()
-        if torch.cuda.is_available():
+        torch.cuda.empty_cache()
            torch.cuda.empty_cache()
    # If custom instance prompts are NOT provided (i.e. the instance prompt is used for all images),
    # pack the statically computed variables appropriately here. This is so that we don't
@@ -1919,11 +1890,6 @@ def main(args):
            text_encoder_lora_layers=text_encoder_lora_layers,
            text_encoder_2_lora_layers=text_encoder_2_lora_layers,
        )
        if args.output_kohya_format:
            lora_state_dict = load_file(f"{args.output_dir}/pytorch_lora_weights.safetensors")
            peft_state_dict = convert_all_state_dict_to_peft(lora_state_dict)
            kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict)
            save_file(kohya_state_dict, f"{args.output_dir}/pytorch_lora_weights_kohya.safetensors")
        # Final inference
        # Load previous pipeline
@@ -53,7 +53,7 @@ from diffusers.utils.torch_utils import is_compiled_module
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -59,7 +59,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -71,7 +71,12 @@ TORCH_DTYPE_MAPPING = {"fp32": torch.float32, "fp16": torch.float16, "bf16": tor
 def log_validation(
-    pipeline, args, accelerator, generator, global_step, is_final_validation=False, enable_autocast=True
+    pipeline,
    args,
    accelerator,
    generator,
    global_step,
    is_final_validation=False,
 ):
    logger.info(
        f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
@@ -91,7 +96,7 @@ def log_validation(
        else Image.open(image_url_or_path).convert("RGB")
    )(args.val_image_url_or_path)
-    with torch.autocast(accelerator.device.type, enabled=enable_autocast):
+    with torch.autocast(str(accelerator.device).replace(":0", ""), enabled=accelerator.mixed_precision == "fp16"):
        edited_images = []
        # Run inference
        for val_img_idx in range(args.num_validation_images):
@@ -492,13 +497,6 @@ def main():
            ),
        )
    logging_dir = os.path.join(args.output_dir, args.logging_dir)
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
@@ -983,13 +981,6 @@ def main():
    if accelerator.is_main_process:
        accelerator.init_trackers("instruct-pix2pix-xl", config=vars(args))
    # Some configurations require autocast to be disabled.
    enable_autocast = True
    if torch.backends.mps.is_available() or (
        accelerator.mixed_precision == "fp16" or accelerator.mixed_precision == "bf16"
    ):
        enable_autocast = False
    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
@@ -1202,7 +1193,6 @@ def main():
                        generator,
                        global_step,
                        is_final_validation=False,
                        enable_autocast=enable_autocast,
                    )
                    if args.use_ema:
@@ -1252,7 +1242,6 @@ def main():
                generator,
                global_step,
                is_final_validation=True,
                enable_autocast=enable_autocast,
            )
    accelerator.end_training()
@@ -52,7 +52,7 @@ if is_wandb_available():
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -46,7 +46,7 @@ from diffusers.utils import check_min_version, is_wandb_available
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -46,7 +46,7 @@ from diffusers.utils import check_min_version, is_wandb_available
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -51,7 +51,7 @@ if is_wandb_available():
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -1,121 +0,0 @@
 This project is an attempt to check if it's possible to apply to [ORPO](https://arxiv.org/abs/2403.07691) on a text-conditioned diffusion model to align it on preference data WITHOUT a reference model. The implementation is based on https://github.com/huggingface/trl/pull/1435/. 
 > [!WARNING] 
 > We assume that MSE in the diffusion formulation approximates the log-probs as required by ORPO (hat-tip to [@kashif](https://github.com/kashif) for the idea). So, please consider this to be extremely experimental.
 ## Training
 Here's training command you can use on a 40GB A100 to validate things on a [small preference
 dataset](https://hf.co/datasets/kashif/pickascore): 
 ```bash
 accelerate launch train_diffusion_orpo_sdxl_lora.py \
  --pretrained_model_name_or_path=stabilityai/stable-diffusion-xl-base-1.0  \
  --pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix \
  --output_dir="diffusion-sdxl-orpo" \
  --mixed_precision="fp16" \
  --dataset_name=kashif/pickascore \
  --train_batch_size=8 \
  --gradient_accumulation_steps=2 \
  --gradient_checkpointing \
  --use_8bit_adam \
  --rank=8 \
  --learning_rate=1e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=2000 \
  --checkpointing_steps=500 \
  --run_validation --validation_steps=50 \
  --seed="0" \
  --report_to="wandb" \
  --push_to_hub
 ```
 We also provide a simple script to scale up the training on the [yuvalkirstain/pickapic_v2](https://huggingface.co/datasets/yuvalkirstain/pickapic_v2) dataset:
 ```bash
 accelerate launch --multi_gpu train_diffusion_orpo_sdxl_lora_wds.py \
  --pretrained_model_name_or_path=stabilityai/stable-diffusion-xl-base-1.0  \
  --pretrained_vae_model_name_or_path=madebyollin/sdxl-vae-fp16-fix \
  --dataset_path="pipe:aws s3 cp s3://diffusion-preference-opt/{00000..00644}.tar -" \
  --output_dir="diffusion-sdxl-orpo-wds" \
  --mixed_precision="fp16" \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --use_8bit_adam \
  --rank=8 \
  --dataloader_num_workers=8 \
  --learning_rate=3e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=50000 \
  --checkpointing_steps=2000 \
  --run_validation --validation_steps=500 \
  --seed="0" \
  --report_to="wandb" \
  --push_to_hub
 ```
 We tested the above on a node of 8 H100s but it should also work on A100s. It requires the `webdataset` library for faster dataloading. Note that we kept the dataset shards on an S3 bucket but it should be also possible to have them stored locally. 
 You can use the code below to convert the original dataset into `webdataset` shards:
 ```python
 import os
 import io
 import ray
 import webdataset as wds
 from datasets import Dataset
 from PIL import Image
 ray.init(num_cpus=8)
 def convert_to_image(im_bytes):
    return Image.open(io.BytesIO(im_bytes)).convert("RGB")
 def main():
    dataset_path = "/pickapic_v2/data"
    wds_shards_path = "/pickapic_v2_webdataset"
    # get all .parquet files in the dataset path
    dataset_files = [
        os.path.join(dataset_path, f)
        for f in os.listdir(dataset_path)
        if f.endswith(".parquet")
    ]
    @ray.remote
    def create_shard(path):
        # get basename of the file
        basename = os.path.basename(path)
        # get the shard number data-00123-of-01034.parquet -> 00123
        shard_num = basename.split("-")[1]
        dataset = Dataset.from_parquet(path)
        # create a webdataset shard
        shard = wds.TarWriter(os.path.join(wds_shards_path, f"{shard_num}.tar"))
        for i, example in enumerate(dataset):
            wds_example = {
                "__key__": str(i),
                "original_prompt.txt": example["caption"],
                "jpg_0.jpg": convert_to_image(example["jpg_0"]),
                "jpg_1.jpg": convert_to_image(example["jpg_1"]),
                "label_0.txt": str(example["label_0"]),
                "label_1.txt": str(example["label_1"])
            }
            shard.write(wds_example)
        shard.close()
    futures = [create_shard.remote(path) for path in dataset_files]
    ray.get(futures)
 if __name__ == "__main__":
    main()
 ```
 ## Inference
 Refer to [sayakpaul/diffusion-sdxl-orpo](https://huggingface.co/sayakpaul/diffusion-sdxl-orpo) for an experimental checkpoint.
@@ -1,7 +0,0 @@
 datasets
 accelerate
 transformers
 torchvision
 wandb
 peft
 webdataset
@@ -23,7 +23,6 @@ TODO:
 6. Integrate to training x
 7. Test
 """
 import copy
 import random
@@ -12,7 +12,7 @@
 # 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.
-"""Conversion script for stable diffusion checkpoints which _only_ contain a controlnet."""
+""" Conversion script for stable diffusion checkpoints which _only_ contain a controlnet. """
 import argparse
 import re
@@ -1005,7 +1005,7 @@ class PromptDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Lo
            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.
+                `._callback_tensor_inputs` attribute of your pipeine class.
        Examples:
@@ -60,7 +60,7 @@ if is_wandb_available():
    import wandb
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -56,7 +56,7 @@ if is_wandb_available():
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -49,7 +49,7 @@ from diffusers.utils import check_min_version
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = logging.getLogger(__name__)
@@ -52,7 +52,7 @@ from diffusers.utils.torch_utils import is_compiled_module
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -64,7 +64,7 @@ from diffusers.utils.torch_utils import is_compiled_module
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -501,12 +501,6 @@ def main(args):
    logging_dir = Path(args.output_dir, args.logging_dir)
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
@@ -979,13 +973,6 @@ def main(args):
    if accelerator.is_main_process:
        accelerator.init_trackers("text2image-fine-tune", config=vars(args))
    # Some configurations require autocast to be disabled.
    enable_autocast = True
    if torch.backends.mps.is_available() or (
        accelerator.mixed_precision == "fp16" or accelerator.mixed_precision == "bf16"
    ):
        enable_autocast = False
    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
@@ -1212,10 +1199,7 @@ def main(args):
                generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
                pipeline_args = {"prompt": args.validation_prompt}
-                with torch.autocast(
+                with torch.cuda.amp.autocast():
                    accelerator.device.type,
                    enabled=enable_autocast,
                ):
                    images = [
                        pipeline(**pipeline_args, generator=generator).images[0]
                        for _ in range(args.num_validation_images)
@@ -54,7 +54,7 @@ from diffusers.utils.torch_utils import is_compiled_module
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -590,12 +590,6 @@ def main(args):
    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
@@ -917,7 +911,6 @@ def main(args):
        )
        precomputed_dataset = precomputed_dataset.with_transform(preprocess_train)
    del compute_vae_encodings_fn, compute_embeddings_fn, text_encoder_one, text_encoder_two
    del text_encoders, tokenizers, vae
    gc.collect()
    torch.cuda.empty_cache()
@@ -986,13 +979,6 @@ def main(args):
        model = model._orig_mod if is_compiled_module(model) else model
        return model
    # Some configurations require autocast to be disabled.
    enable_autocast = True
    if torch.backends.mps.is_available() or (
        accelerator.mixed_precision == "fp16" or accelerator.mixed_precision == "bf16"
    ):
        enable_autocast = False
    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
@@ -1226,10 +1212,7 @@ def main(args):
                generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
                pipeline_args = {"prompt": args.validation_prompt}
-                with torch.autocast(
+                with torch.cuda.amp.autocast():
                    accelerator.device.type,
                    enabled=enable_autocast,
                ):
                    images = [
                        pipeline(**pipeline_args, generator=generator, num_inference_steps=25).images[0]
                        for _ in range(args.num_validation_images)
@@ -1284,7 +1267,7 @@ def main(args):
        if args.validation_prompt and args.num_validation_images > 0:
            pipeline = pipeline.to(accelerator.device)
            generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
-            with torch.autocast(accelerator.device.type, enabled=enable_autocast):
+            with torch.cuda.amp.autocast():
                images = [
                    pipeline(args.validation_prompt, num_inference_steps=25, generator=generator).images[0]
                    for _ in range(args.num_validation_images)
@@ -80,7 +80,7 @@ else:
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -56,7 +56,7 @@ else:
 # ------------------------------------------------------------------------------
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = logging.getLogger(__name__)
@@ -76,7 +76,7 @@ else:
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__)
@@ -29,7 +29,7 @@ from diffusers.utils.import_utils import is_xformers_available
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -50,7 +50,7 @@ if is_wandb_available():
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -51,7 +51,7 @@ if is_wandb_available():
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.28.0.dev0")
+check_min_version("0.27.0.dev0")
 logger = get_logger(__name__, log_level="INFO")
@@ -12,7 +12,7 @@
 # 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.
-"""Conversion script for the LDM checkpoints."""
+""" Conversion script for the LDM checkpoints. """
 import argparse
 import json
@@ -12,7 +12,7 @@
 # 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.
-"""Conversion script for the LDM checkpoints."""
+""" Conversion script for the LDM checkpoints. """
 import argparse
@@ -1195,9 +1195,9 @@ def superres_check_against_original(dump_path, unet_checkpoint_path):
        if_II_model = IFStageIII(device="cuda", dir_or_name=orig_path, model_kwargs={"precision": "fp32"}).model
    batch_size = 1
-    channels = model.config.in_channels // 2
+    channels = model.in_channels // 2
-    height = model.config.sample_size
+    height = model.sample_size
-    width = model.config.sample_size
+    width = model.sample_size
    height = 1024
    width = 1024
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Sayak Paul	47ee2a737a	Merge branch 'main' into controlnet-test-fixes	2024-03-15 12:30:25 +05:30
Dhruv Nair	94fc2d3fe6	update	2024-03-14 06:53:52 +00:00
Dhruv Nair	503e359204	update	2024-03-14 06:47:29 +00:00